ABC解説

ABC369 C問題の解説（Python）

はじめに

この記事では、ABC369のC問題を解説していきます。

ABC（AtCoder Beginner Contest）とは、AtCoderが開催している、競技プログラミングコンテストです。

ABC369　C – Count Arithmetic Subarrays

問題

問題

長さNの整数列 A=(A1,…,AN) が与えられます。

1≦l≦r≦N を満たす整数の組 (l,r) で、数列 (Al,Al+1,…,Ar) が等差数列となるようなものが何通りあるかを求める問題です。

制約

  • 1≦N≦2×105
  • 1≦Ai≦109

思考の筋道

入力例1で考えてみましょう。

入力例1

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 となります。

一般に同じ差がn連続したとき、その区間で等差数列は
1+…+n=12n(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)

以上、ABC369のC問題の解説でした！

では、またね。

リンク

ABC369
情報科学の部屋へ