Closed1
ABC369Cmemo
長さ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
を作成し、ランレングス圧縮を行った後、各ブロックに対して条件を満たす部分列をカウント
フロー
-
差分配列の作成:
- 数列
A
の隣接する要素の差を計算し、差分配列B
を作成
- 数列
-
ランレングス圧縮:
- 差分配列
B
をランレングス圧縮し、連続して同じ値を持つ部分列をf(i)
として表現
- 差分配列
-
部分列の数を計算:
- 各ブロック
f(i)
に対して、その中で条件を満たす区間の数をD_i * (D_i + 1) / 2
で計算し、それらを合計
- 各ブロック
このスクラップは2024/09/01にクローズされました