目次
はじめに
AtCoder Beginner Contest（略称：ABC）とは、AtCoderが開催している、競技プログラミングコンテストのことです。
今回はABC369のC問題を解説していきます。
ABC369 C – Count Arithmetic Subarrays
問題
長さNの整数列 \(A=(A_1, …, A_N )\) が与えられます。
\(1 \leqq l \leqq r \leqq N \) を満たす整数の組 \((l, r)\) で、数列 \((A_l, A_{l+1},…, A_r )\) が等差数列となるようなものが何通りあるかを求める問題です。
制約
- \(1 \leqq N \leqq 2×10^5\)
- \(1 \leqq A_i \leqq 10^9\)
思考の筋道
入力例1で考えてみましょう。
入力例1
一般に同じ差がn連続したとき、その区間で等差数列は
\(1+…+n=\displaystyle\frac{1}{2}n(n+1)\) 個作れます。
これにより、countを計算できますね。
下のコードでは、rensaという変数を使って、同じ差が何連続しているかを管理しています。
コード
N = int(input())
A = list(map(int, input().split()))
B = [0] * (N - 1)
for i in range(N - 1):
B[i] = A[i + 1] - A[i]
count = N
rensa = 1
for i in range(N - 1):
if i == N - 2:
count += rensa * (rensa + 1) // 2
elif B[i + 1] == B[i]:
rensa += 1
else:
count += rensa * (rensa + 1) // 2
rensa = 1
print(count)
\(A=(3, 6, 9, 3 )\)
まず、長さ1の数列はすべて等差数列であることから、countの初期値をN（今回の例だと4）にしておきます。
今回、興味があるのは隣接する項の差なので、隣接する項の差の数列（階差数列）B を作ります。
\(B=(3, 3, -6 )\)
3が2連続であることから、この区間で長さ3の等差数列1つと長さ2の等差数列2つ、つまり、計3つの等差数列が作れます。
長さ3の等差数列：(3, 6, 9)のこと
長さ2の等差数列：(3, 6)と(6, 9)のこと
また-6から、この区間で長さ2の等差数列1つが作れます。
長さ2の等差数列：(9, 3)のこと
以上より、count = 4 + 3 + 1 = 8 となります。