u++の備忘録

「累積和を何も考えずに書けるようにする！」をPythonで解く

python 競プロ

「累積和を何も考えずに書けるようにする！」という記事に記載されていた計5問をPythonで解きました。Pythonだと、itertools.accumulateを利用することで1次元累積和が簡単に記述できます。


累積和の例題

問題 1： AOJ 0516 - 最大の和 (JOI 2006 本選 A)

from itertools import accumulate


N, K = list(map(int, input().split()))
A = list(map(int, input().split()))

# 累積和
B = [0] + A
B = list(accumulate(B))

ans = []
for i in range(N-3):
    ans.append(B[i+3]-B[i])

print(max(ans))

問題 2： AtCoder ABC 084 D - 2017-like Number

import math
from itertools import accumulate


def is_prime(n):
    if n == 1:
        return 0

    for k in range(2, int(math.sqrt(n)) + 1):
        if n % k == 0:
            return 0

    return 1


Q = int(input())
L = [0]*Q
R = [0]*Q

for i in range(Q):
    L[i], R[i] = list(map(int, input().split()))

min_L = min(L)
max_R = max(R)

check = []
for i in range(min_L, max_R+1):
    if (i % 2 == 1):
        flg = (is_prime(i) and is_prime((i+1)/2))
    else:
        flg = 0
    check.append(flg)

# 累積和
check = [0] + check
check = list(accumulate(check))

for i in range(Q):
    print(check[R[i] - min_L + 1] - check[L[i] - min_L])

エラトステネスのふるいを利用する場合。

from itertools import accumulate


# エラトステネスの篩
# is_prime := 1~nが素数か否か
# table := 2~nのうち素数
def sieve(n):
    is_prime = [True for _ in range(n)]
    is_prime[0] = False

    for i in range(2, n+1):
        if is_prime[i-1]:
            j = 2 * i
            while j <= n:
                is_prime[j-1] = False
                j += i
    table = [i for i in range(1, n+1) if is_prime[i-1]]
    return is_prime, table


Q = int(input())
L = [0]*Q
R = [0]*Q

for i in range(Q):
    L[i], R[i] = list(map(int, input().split()))

max_R = max(R)
is_prime, table = sieve(max_R)

check = [0]*(max_R + 1)
for i in range(0, max_R + 1):
    if (i % 2 == 1):
        if (is_prime[i-1] and is_prime[(i+1)//2-1]):
            check[i] = 1

# 累積和
check = [0] + check
check = list(accumulate(check))

for i in range(Q):
    print(check[R[i] + 1] - check[L[i]])

問題 3： AtCoder ABC 122 C - GeT AC

from itertools import accumulate


N, Q = list(map(int, input().split()))
S = input()
A = [list(map(int, input().split())) for i in range(Q)]

flg = []
for i in range(N-1):
    tmp_S = S[i:i+2]
    if tmp_S == 'AC':
        flg.append(1)
    else:
        flg.append(0)

flg2 = [0] + flg
flg2 = list(accumulate(flg2))

for a in A:
    print(flg2[a[1]-1] - flg2[a[0]-1])

問題 4： AtCoder AGC 023 A - Zero-Sum Ranges

from itertools import accumulate
import collections


N = int(input())
A = list(map(int, input().split()))

# 累積和
B = [0] + A
B = list(accumulate(B))

# 個数を数える
c = collections.Counter(B)
cm = c.most_common()

cnt = 0
for v in cm:
    if v[1] > 1:
        cnt += (v[1]*(v[1]-1)//2)

print(cnt)

二次元累積和

問題 5: AtCoder ABC 005 D - おいしいたこ焼きの焼き方

N = int(input())
D = [list(map(int, input().split())) for i in range(N)]
Q = int(input())
P = [int(input()) for i in range(Q)]

# 2次元累積和
S = [[0 for i in range(N+1)] for j in range(N+1)]
for i in range(N):
    for j in range(N):
        S[i+1][j+1] = S[i][j+1] + S[i+1][j] - S[i][j] + D[i][j]

val = [0]*(N*N+1)   # val[v] := 面積が v の長方形領域の総和の最大値
for x1 in range(0, N):
    for x2 in range(x1+1, N+1):
        for y1 in range(0, N):
            for y2 in range(y1+1, N+1):
                area = (x2-x1)*(y2-y1)
                v = S[x2][y2] - S[x1][y2] - S[x2][y1] + S[x1][y1]
                val[area] = max(val[area], v)

# val[v] := 面積が v 「以下」の長方形領域の総和の最大値
for i in range(0, N*N):
    val[i+1] = max(val[i+1], val[i])

for p in P:
    print(val[p])