ABC355のA~D, F問題をpythonで解説していきます。筆者はA-Dの4完でした。

A - Who Ate the Cake?

問題

高橋君のケーキが誰かに食べられてしまいました。ケーキを食べた犯人の候補として、人 $1$、人 $2$、人 $3$ の三人が挙げられています。

犯人の目撃者はりんごさんとすぬけくんの二人がいます。りんごさんは人 $A$ が犯人でないことを覚えており、すぬけくんは人 $B$ が犯人でないことを覚えています。

二人の記憶から犯人を一人に特定できるかどうか判定し、特定できるならばその人の番号を出力してください。

考察

人1から人3まで順に、無実リストに入っているかチェックします。入っていない人をリストに入れていき、最終的にその長さが1より大きかったら怪しい人が複数いるため特定不可の-1、そうでなればそのリストに入っている人を答えとして出力すればよいです。

コード

A.py

ans = []
candi = list(map(int, input().split()))
for i in range(1, 4):
    if i not in candi:
        ans.append(i)
if len(ans) > 1:
    print(-1)
else:
    print(ans[0])

B - Piano 2

問題

長さ $N$ の数列 $A=(A_1, A_2, \dots, A_N)$ と、長さ $M$ の数列 $B=(B_1, B_2, \dots, B_M)$ が与えられます。ここで、$A, B$ のすべての要素は互いに相異なります。$A, B$ のすべての要素を昇順に並べた長さ $N+M$ の数列 $C=(C_1, C_2, \dots, C_{N+M})$ において、$A$ に現れる要素が2つ連続するかどうか判定してください。

考察

Aの要素は0、Bの要素は1のようにラベル付けをし、それらラベル付けされた要素のあるリストを結合してソートすれば、あとはラベル0が2個連続で並んでいるかどうかをチェックすればよいです。並んでいれば"Yes"、そうでなければ"No"と出力します。

コード

B.py

N, M = map(int, input().split())
A = list(map(int, input().split()))
B = list(map(int, input().split()))
A = [(a, 0) for a in A]
B = [(b, 1) for b in B]
C = A + B
C.sort()
for i in range(len(C) - 1):
    (_, d1), (_, d2) = C[i], C[i + 1]
    if d1 == d2 == 0:
        print("Yes")
        exit()
print("No")

C - Bingo 2

問題

縦 $N$ 行、横 $N$ 列のマス目があり、上から $i$ 行目、左から $j$ 列目のマスには整数 $N \times (i-1) + j$ が書かれています。

今から $T$ ターンにわたって相異なる整数が宣言されます。$i$ ターン目には $A_i$ が宣言され、$A_i$ が書かれたマスに印をつけます。初めてビンゴを達成するのは何ターン目か求めてください。ただし、$T$ ターンの中でビンゴを達成しない場合は -1 を出力してください。

ここで、ビンゴを達成するとは以下のいずれかのうち少なくとも一つ満たされることを言います。

マス目の横の列であって、列に含まれる $N$ 個のマスすべてに印がついているものが存在する
マス目の縦の列であって、列に含まれる $N$ 個のマスすべてに印がついているものが存在する
マス目の対角線の列であって、列に含まれる $N$ 個のマスすべてに印がついているものが存在する

考察

それぞれの縦横斜めの列に対して、印のついていないマスの個数を管理します。あとは毎ターン値$A$に対応する列の値を1減らし、どれかの列の値が0となった時点でその列がビンゴしたと判定できます。

コード

C.py

N, T = map(int, input().split())
tate = [N] * N
yoko = [N] * N
naname = [N] * 2
A = list(map(int, input().split()))
for t, a in enumerate(A):
    a -= 1
    i, j = divmod(a, N)
    tate[i] -= 1
    yoko[j] -= 1
    if i == j:
        naname[0] -= 1
    if j == N - i - 1:
        naname[1] -= 1
    if naname[0] == 0 or naname[1] == 0 or tate[i] == 0 or yoko[j] == 0:
        print(t + 1)
        exit()
print(-1)

D - Intersecting Intervals

問題

$N$ 個の実数の区間が与えられます。$i,(1 \leq i \leq N)$ 番目の区間は $[l_i, r_i]$ です。$i$ 番目の区間と $j$ 番目の区間が共通部分を持つような組 $(i, j),(1 \leq i < j \leq N)$ の個数を求めてください。

考察

rの降順、lの昇順でソートしたとき、ある区間$[l_i, r_i]$が区間$[l_j, r_j]$と共通部分をもつかどうかは、$r_i < l_j$かどうかでチェックできます($r_i<r_j$とする)。$l_j$に対応する部分を+1しておけば、ある区間$[l_i, r_i]$が共通部分をもたない区間の個数は、$r_i$に対応する部分より後ろの和で求められるので、1点区間更新(+1する)と区間クエリ($r_i$以降の和)を高速に行えるデータ構造であるセグメント木で情報を管理すればよいです。(ただし、今回は$l$や$r$が$1$から$10^9$の値をとり得るので、座標圧縮が必要な点に注意が必要です。)

コード

D.py

# https://raw.githubusercontent.com/shakayami/ACL-for-python/master/segtree.py
class segtree:
    n = 1
    size = 1
    log = 2
    d = [0]
    op = None
    e = 10**15

    def __init__(self, V, OP, E):
        self.n = len(V)
        self.op = OP
        self.e = E
        self.log = (self.n - 1).bit_length()
        self.size = 1 << self.log
        self.d = [E for i in range(2 * self.size)]
        for i in range(self.n):
            self.d[self.size + i] = V[i]
        for i in range(self.size - 1, 0, -1):
            self.update(i)

    def set(self, p, x):
        assert 0 <= p and p < self.n
        p += self.size
        self.d[p] = x
        for i in range(1, self.log + 1):
            self.update(p >> i)

    def get(self, p):
        assert 0 <= p and p < self.n
        return self.d[p + self.size]

    def prod(self, l, r):
        assert 0 <= l and l <= r and r <= self.n
        sml = self.e
        smr = self.e
        l += self.size
        r += self.size
        while l < r:
            if l & 1:
                sml = self.op(sml, self.d[l])
                l += 1
            if r & 1:
                smr = self.op(self.d[r - 1], smr)
                r -= 1
            l >>= 1
            r >>= 1
        return self.op(sml, smr)

    def all_prod(self):
        return self.d[1]

    def max_right(self, l, f):
        assert 0 <= l and l <= self.n
        assert f(self.e)
        if l == self.n:
            return self.n
        l += self.size
        sm = self.e
        while 1:
            while l % 2 == 0:
                l >>= 1
            if not (f(self.op(sm, self.d[l]))):
                while l < self.size:
                    l = 2 * l
                    if f(self.op(sm, self.d[l])):
                        sm = self.op(sm, self.d[l])
                        l += 1
                return l - self.size
            sm = self.op(sm, self.d[l])
            l += 1
            if (l & -l) == l:
                break
        return self.n

    def min_left(self, r, f):
        assert 0 <= r and r <= self.n
        assert f(self.e)
        if r == 0:
            return 0
        r += self.size
        sm = self.e
        while 1:
            r -= 1
            while r > 1 and (r % 2):
                r >>= 1
            if not (f(self.op(self.d[r], sm))):
                while r < self.size:
                    r = 2 * r + 1
                    if f(self.op(self.d[r], sm)):
                        sm = self.op(self.d[r], sm)
                        r -= 1
                return r + 1 - self.size
            sm = self.op(self.d[r], sm)
            if (r & -r) == r:
                break
        return 0

    def update(self, k):
        self.d[k] = self.op(self.d[2 * k], self.d[2 * k + 1])

    def __str__(self):
        return str([self.get(i) for i in range(self.n)])

    def get_list(self):
        return [self.get(i) for i in range(self.n)]  # オリジナルで追加


# RSQ
def op(x, y):
    return x + y


N = int(input())
LR = [list(map(int, input().split())) for _ in range(N)]
X = []
for l, r in LR:
    X.append(l)
    X.append(r)


def compress(arr):  # 座標圧縮
    (*XS,) = set(arr)
    XS.sort()
    return {cmp_e: cmp_i for cmp_i, cmp_e in enumerate(XS)}


d = compress(X)
LR = [(-d[r], d[l]) for l, r in LR]  # r降順, l昇順
LR.sort()
L = len(d) + 1
seg = segtree([0] * L, op, 0)
ans = 0
for i, (r, l) in enumerate(LR):
    r *= -1
    minus = seg.prod(r + 1, L)
    ans += i - minus
    seg.set(l, seg.get(l) + 1)
print(ans)

F - MST Query

問題

頂点に $1$ から $N$ の番号が、辺に $1$ から $N-1$ の番号が付いた $N$ 頂点 $N-1$ 辺の重み付き無向連結グラフ $G$ が与えられます。辺 $i$ は頂点 $a_i$ と頂点 $b_i$ を結んでおり、その重みは $c_i$ です。

$Q$ 個のクエリが与えられるので順に処理してください。$i$ 番目のクエリは以下で説明されます。

整数 $u_i, v_i, w_i$ が与えられる。$G$ の頂点 $u_i, v_i$ の間に重み $w_i$ の辺を追加する。その後、$G$ の最小全域木に含まれる辺の重みの和を出力する。

考察

$uf[i]: 重みi以下の辺のみで構成されたUnionFind$
とします。
ある辺(u, v, w)によってMSTの重み和がどれくらい減少するかは、現在頂点u, vを連結にするために最小でどれだけの重みが使われているかわかればよく、それはuf[j]においてu, vが連結となっているという条件を満たすような最小のjに対応することが考えるとわかります。よってこの辺の追加により、$max(0, j - w)$だけ重みが減少するとわかります。

コード

F.py

class UnionFind:
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

        self.parents[x] += self.parents[y]
        self.parents[y] = x

    def size(self, x):
        return -self.parents[self.find(x)]

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def members(self, x):  # 多用すると重い
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        return len(self.roots())

    def all_group_members(self):
        group_members = defaultdict(list)
        for member in range(self.n):
            group_members[self.find(member)].append(member)
        return group_members

    def __str__(self):
        return "\n".join(f"{r}: {m}" for r, m in self.all_group_members().items())


N, Q = map(int, input().split())
uf = [UnionFind(N + 1) for _ in range(11)]
now = 0
for _ in range(N - 1):
    a, b, c = map(int, input().split())
    now += c
    for w in range(c, 11):
        uf[w].union(a, b)
for _ in range(Q):
    u, v, w = map(int, input().split())
    for j in range(11):
        if uf[j].same(u, v):
            break
        if j >= w:
            uf[j].union(u, v)
    now -= max(0, j - w)
    print(now)