aacord’s memo

abcを中心にpythonで解いた問題のメモ、整理をしています。緑になった。

AGC 010 B - Boxes (python)

数列から等差数列を引いていって [0]*n にできるかを答える問題
n = 5 の場合、引く候補は [1,2,3,4,5], [2,3,4,5,1], [3,4,5,1,2], [4,5,1,2,3], [5,1,2,3,4] の 5 つあるが、
どれも総和は 15 であり、引ける回数 k は sum(a) // 15 になる
引いた後の遷移について考えると、等差数列を引くときは元の数列の隣り合う差を考えると上手くいくことがある。
D[i] = A[i+1] - A[i] とすると、D[i] は一回等差数列を引くごとに、1減るか、n-1増える。
よって k 回操作を行い m 回 n-1 増える場合、D[i] = D[i] - k + m*n となる
よってこの問題の必要十分条件は、全ての D[i] について D[i] - k % n = 0 かつ D[i] - k ≤ 0 かつ Σ m ≤ k となる

n = int(input())
a = list(map(int,input().split()))
d = []
for i in range(n-1):
  d.append(a[i+1]-a[i])
  
f = True
if sum(a) % (n*(n+1)//2) != 0:
  f = False
  k = 0
else:
  k = sum(a)//(n*(n+1)//2)
  
cnt = 0
for i in d:
  j = i-k
  if j > 0 or j % n != 0:
    f = False
  else:
    cnt -= j//n
    
if cnt > k:
  f = False
  
if f:
  print('YES')
else:
  print('NO')