ABC369 C問題（Count Arithmetic Subarrays）を解く

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

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

C問題 Count Arithmetic Subarrays（Difficulty : 323）

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

ABC369 C問題 Count Arithmetic Subarrays

数列の差分を計算して、この差分を処理します。AtCoder Problems による Difficulty は 323 でした。

解答案

与えられた数列の隣の要素との差分を新しい数列として処理します。この数列は、与えられた数列の要素から1個少なくなります。

入力例1の数列：3, 6, 9, 3
隣の要素との差分の数列：3, 3, -6

同じ差分が続きます。2回目の3で、等差数列は2個増えます（3→6→9と6→9）。同じ差分が続けば、その連続した回数だけ等差数列の数が増えます。長さが1の数列は、この方法ではカウントできません。長さ数列の個数は $N$ 個あるため、別に解答に加えます。

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

差分の数列 d[0] は、あり得ない値（INF）としました（17行目）。後は、d[i-1]とd[i]を比較して、同じであれば、連続した回数（conseq）をひとつ増やします。異なれば、連続した回数を1にします。どちらの場合も解答（result）に conseq を加えます。

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

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

typedef long long int ll;
#define INF (1LL << 60)

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

	vector<ll> d(n);
	d[0] = INF;
	for (int i = 1; i < n; ++i) {
		d[i] = a[i] - a[i - 1];
	}

	ll result = n;
	int conseq = 1;
	for (int i = 1; i < n; ++i) {
		if (d[i - 1] == d[i]) {
			++conseq;
		} else {
			conseq = 1;
		}
		result += conseq;
	}

	cout << result << endl;

	return 0;
}

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

Python プログラム例（ABC369C）

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

"""AtCoder Beginner Contest 369 C"""
INF = 1 << 60
n = int(input())
a = list(map(int, input().split()))
d = [0] * n
d[0] = INF
for i in range(1, n):
    d[i] = a[i] - a[i - 1]

result = n
conseq = 1
for i in range(1, n):
    if d[i - 1] == d[i]:
        conseq += 1
    else:
        conseq = 1
    result += conseq

print(result)

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

最後に

今回のコンテストは、E問題を解くことができましたが、このC問題は後回しにして解くことができませんでした。想定していた解法は正しかったため残念です。

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