結果

問題 No.2956 Substitute with Average
ユーザー dp_ijk
提出日時 2024-11-09 00:00:13
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 962 ms / 3,000 ms
コード長 906 bytes
コンパイル時間 333 ms
コンパイル使用メモリ 82,404 KB
実行使用メモリ 237,176 KB
最終ジャッジ日時 2024-11-09 00:00:31
合計ジャッジ時間 17,803 ms
ジャッジサーバーID
(参考情報)
judge3 / judge2
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 3
other AC * 25
権限があれば一括ダウンロードができます

ソースコード

diff #

from collections import defaultdict

from itertools import accumulate
N = int(input())
A = list(map(int, input().split()))


def helper(A):
   N = len(A)
   W = list(accumulate(A, initial=0))
   dp = [0]*(30+1)
   for x in range(1, 31):
      d = defaultdict(int)
      term = 0
      for i in range(N):
         t = W[i+1] - (i+1) * x
         if A[i] == x:
            term += d[t]
         s = W[i] - i*x
         d[s] += 1
      dp[x] += term
   return sum(dp)


def resolve(A):
   N = len(A)
   W = list(accumulate(A, initial=0))
   dp = [0]*(30+1)
   for x in range(1, 31):
      d = defaultdict(int)
      term = 0
      for i in range(N):
         if A[i] != x: continue
         t = W[i+1] - (i+1) * x
         term += d[t]
         s = W[i] - i*x
         d[s] += 1
      dp[x] += term
   return sum(dp)


ans = N*(N-1) // 2
ans -= helper(A)
ans -= helper(A[::-1])
ans += resolve(A)
print(ans+1)
0