ABC367 D問題（Pedometer）を解く

AtCoder が提供しているABC（AtCoder Beginner Contest）367 D問題をC++とPythonで解いてみました。ABC367は、2024年8月17日21:00に実施されました。

AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。

D問題 Pedometer（Difficulty : 1037）

問題の詳細は、リンク先をご覧ください。

ABC367 D問題 Pedometer

累積和の剰余を求め、その結果を用いて処理します。AtCoder Problems による Difficulty は 1037 でした。

解答案

入力例1を表にしてみます。説明のため2周分のデータを記します。

時計回りに1周未満の移動で剰余が同じとなる組み合わせを探します。具体的には、以下の4通りとなります。

• 1 → 3（剰余0で同じ）
• 3 → 2（剰余0で同じ）
• 4 → 1（剰余1で同じ）
• 4 → 3（剰余1で同じ）

これを上手くカウントするためには、以下のように考えます。

• 2回目の1に注目する。その前の2から4までで剰余が同じになるものは1か所（4→1）
• 2回目の2に注目する。その前の3から1までで剰余が同じになるものは1か所（3→2）
• 2回目の3に注目する。その前の4から2までで剰余が同じになるものは2か所（4→3と1→3）
• 2回目の4に注目する。その前の1から3までで剰余が同じになるものはない。

2周分の剰余を求めておいて、注目している数と同じ剰余をもつ3個前から直前までの個数を加えます。

C++ プログラム例（ABC367D）

上の考察をプログラムとして表現しました。

• 配列aを読み込み、2周分に拡張します（14行目の処理）。
• 2周分の累積和の剰余を配列 sum として計算します（17ー22行目）。
• 1周分の剰余の頻度を求めておきます（24ー27行目）。
• 一番中心となる処理が、31ー33行目です。
  • n個前の剰余を頻度から抜きます（31行目）。
  • 2周目で注目している数と同じ剰余の頻度を加えます（32行目）。
  • その剰余の頻度をひとつ増やします（33行目）。

以下が、C++プログラムです。

#include <bits/stdc++.h>
using namespace std;

typedef long long int ll;

int main()
{
	int n;
	ll m;
	cin >> n >> m;
	vector<ll> a(2 * n);
	for (int i = 0; i < n; ++i) {
		cin >> a[i];
		a[n + i] = a[i];
	}

	vector<ll> sum(2 * n, 0);
	sum[0] = 0;
	for (int i = 0; i < 2 * n - 1; ++i) {
		sum[i + 1] = sum[i] + a[i];
		sum[i + 1] %= m;
	}

	map<int, int> mod;
	for (int i = 0; i < n; ++i) {
		++mod[sum[i]];
	}

	ll result = 0;
	for (int i = n; i < 2 * n; ++i) {
		--mod[sum[i - n]];
		result += mod[sum[i]];
		++mod[sum[i]];
	}

	cout << result << endl;

	return 0;
}

AC（Accepted＝正しいプログラム）と判定されました。

Python プログラム例（ABC367D）

Python版も基本的な考え方は同じです。以下がプログラムです。

"""AtCoder Beginner Contest 367 D"""
n, m = map(int, input().split())
a = list(map(int, input().split()))
for i in range(n):
    a.append(a[i])

s = [0] * (2 * n)
s[0] = 0
for i in range(2 * n - 1):
    s[i + 1] = s[i] + a[i]
    s[i + 1] %= m

mod = [0] * m
for i in range(n):
    mod[s[i]] += 1

result = 0
for i in range(n, 2 * n):
    mod[s[i - n]] -= 1
    result += mod[s[i]]
    mod[s[i]] += 1

print(result)

こちらも「AC」と判定されました。

最後に

基本的なアイデアを思いついていましたが、プログラムとして実装できませんでした。公式解説のプログラムは、すっきりと書けていて勉強になりました。

引き続き ABC の問題を紹介していきます。