ABC340をPythonで（A～E）

鹿島建設プログラミングコンテスト2024（AtCoder Beginner Contest 340）の解答等のまとめ

A問題

range関数は3つパラメータを入れてると、1つ目から２つ目の手前まで３つ目の間隔で値を出力してくれる

A

a, b, d = map(int, input().split())
print(*[i for i in range(a, b + 1, d)])

B問題

問題群の通りにやる

B

A = []
for _ in range(int(input())):
    com, x = map(int, input().split())
    if com == 1:
        A.append(x)
    else:
        print(A[-x])

C問題

メモ化再帰
lru_cache　を思い出せなかったから辞書を使った

n = int(input())

memo = dict()
memo[0] = memo[1] = 0


def calc(x):
    global memo
    if x in memo:
        return memo[x]
    else:
        floor, ceil = x // 2, x - x // 2
        res = calc(floor) + calc(ceil) + x
        memo[x] = res
        return res


print(calc(n))

D問題

ダイクストラ

D

from heapq import heappop, heappush

n = int(input())

edge = [[] for _ in range(n + 1)]
for i in range(n - 1):
    a, b, x = map(int, input().split())
    edge[i].append([i + 1, a])
    edge[i].append([x - 1, b])

q = [[0, 0]]
INF = float("inf")
dist = [INF] * n
dist[0] = 0
while q:
    d_i, now = heappop(q)
    if dist[now] < d_i:
        continue
    for to, d_j in edge[now]:
        if dist[to] > dist[now] + d_j:
            dist[to] = dist[now] + d_j
            heappush(q, [dist[to], to])

print(dist[-1])

E問題

RangeAddQuery
おそらく遅延セグ木でRAQを処理するものだと思ったがなんやかんやでBITで出来た

$i$番目の和を見ると$a_i$となるように埋め込む

$b_i$のクエリの時は

1. $b_i$の値$x$を取得する
2. $b_{i+1}$に$x$を加え、$b_i$を$0$に
3. 全体に$x //n $を加える
4. １個多くもらうところに１加わるように頑張る

E

class BinaryIndexTree:
    def __init__(self, n):
        self.size = n
        self.tree = [0] * (n + 1)

    def sum(self, i):
        s = 0
        while i > 0:
            s += self.tree[i]
            i -= i & -i
        return s

    def add(self, i, x):
        while i <= self.size:
            self.tree[i] += x
            i += i & -i

    def range_sum(self, i, j):
        if i == 0:
            return self.sum(j)
        return self.sum(j) - self.sum(i - 1)


n, m = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(lambda x:int(x) + 1, input().split()))

BIT = BinaryIndexTree(n * 2 + 10)
s = 0
for i in range(n):
    BIT.add(i + 1, a[i] - s)
    s = a[i]

for b_i in b:
    x = BIT.sum(b_i)
    BIT.add(b_i, -x)
    BIT.add(b_i + 1, x)
    if x // n > 0:
        BIT.add(1, x // n)
    
    if x % n > 0:
        BIT.add(b_i + 1, 1)
        BIT.add(b_i + x % n + 1, -1)
        if b_i + x % n > n:
            BIT.add(1, 1)
            BIT.add(b_i + x % n - n + 1, -1)

ans = [BIT.sum(i) for i in range(1, n + 1)]
print(*ans)