Closed1

ABC369Cmemo

junjun

長さNの数列がある A=(A_1...A_N)
1≦l≦r≦Nを満たす組(l,r)
数列(A_l,A_l+1...A_r)が等差数列である組み合わせは何通りかを求めて
入力形式
N
A_1 A_2...A_N

コード

def count_arithmetic_subarrays(A):
    N = len(A)
    if N < 2:
        return N

    # 1. 差分配列 B を作成
    B = [A[i+1] - A[i] for i in range(N-1)]

    # 2. ランレングス圧縮
    f = []
    current_diff = B[0]
    length = 1

    for i in range(1, len(B)):
        if B[i] == current_diff:
            length += 1
        else:
            f.append(length)
            current_diff = B[i]
            length = 1
    f.append(length)

    # 3. 部分列の数を計算
    total_count = 0
    for D_i in f:
        total_count += D_i * (D_i + 1) // 2

    # 4. 各要素単体も等差数列としてカウント
    total_count += N

    return total_count

# 入力
N = int(input("N: "))  # 数列の長さ
A = list(map(int, input("A: ").split()))  # 数列

# 計算
count = count_arithmetic_subarrays(A)

# 出力
print(count)

数列 A から差分配列 B を作成し、ランレングス圧縮を行った後、各ブロックに対して条件を満たす部分列をカウント

フロー

  1. 差分配列の作成:

    • 数列 A の隣接する要素の差を計算し、差分配列 B を作成
  2. ランレングス圧縮:

    • 差分配列 B をランレングス圧縮し、連続して同じ値を持つ部分列を f(i) として表現
  3. 部分列の数を計算:

    • 各ブロック f(i) に対して、その中で条件を満たす区間の数を D_i * (D_i + 1) / 2 で計算し、それらを合計
このスクラップは2024/09/01にクローズされました