AtCoder ABC420 解いてみた(PythonでABCDE問題)

ABC420を解いてみました

ABC420、リアルタイム参加は出来なかったので、問題だけ解いてみました。
ABCDEまでの5問分の解答を載せています。

A - What month is it?

月を足して12を超えたら、12を引くという問題です。
冷静に考えると (X + Y - 1) % 12 + 1 が綺麗な解法なんですが、この問題の条件なら下のコードでも大丈夫です...。

X, Y = map(int, input().split())
ans = X + Y 
if ans > 12:
    ans -= 12
print(ans)

B - Most Minority

Bにしては、問題文がややこしい印象の問題でした。
M回の投票について、地道に各人のスコアを計算していきます。

N, M = map(int, input().split())
S = []
for i in range(N):
    S.append(input())

score = [0] * N # 人1から人Nのスコアを格納するリスト

# M回の投票を行う
for j in range(M):
    # "0"と"1"の数を数える
    count_zero = 0
    count_one = 0
    for i in range(N):
        if S[i][j] == "0":
            count_zero += 1
        else:
            count_one += 1

    # 0の方が多数？を判定
    is_zero_better = count_zero > count_one

    for i in range(N):
        if is_zero_better and S[i][j] == "1":
            # 0が多数なら、1を選んでる人にスコアを+1
            score[i] += 1
        elif (not is_zero_better) and S[i][j] == "0":
            # 1が多数なら、0を選んでる人にスコアを+1
            score[i] += 1

# 最大スコアを取った人を全員出力
max_score = max(score)
answer = []
for i in range(N):
    if score[i] == max_score:
        answer.append(i + 1)

print(*answer)

C - Sum of Min Query

クエリのたびにAとBの最小値を求めると、計算量が多くなってしまいます。
そこで、AとBの各要素で小さい方を並べたリストを用意します(min_list)。
また、min_list の総和を sum_count として持っておきます。

こうすることで、クエリのたびに min_list と sum_count を更新するだけで、簡単にA・B小さいの方の総和を求めることができます。

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

# AとBの各要素で小さい方を並べたリストを用意
min_list = [min(a, b) for a, b in zip(A, B)]

# min_listの総和を計算しておく
sum_count = sum(min_list)

# クエリを処理していく
for i in range(Q):
    c, x, v = input().split()
    x = int(x)
    v = int(v)

    # AかBのどちらかを更新
    if c == "A":
        A[x - 1] = v
    else:
        B[x - 1] = v

    # min_listとsum_countを更新
    before = min_list[x - 1]
    after = min(A[x - 1], B[x - 1])
    if after < before:
        sum_count -= before - after
        min_list[x - 1] = after
    elif after > before:
        sum_count += after - before
        min_list[x - 1] = after

    # 総和を出力
    print(sum_count)

D - Toggle Maze

最短でゴールまでたどり着く操作回数を求める迷路問題...ということで、幅優先探索BFSで解きます。

注意すべきはスイッチマスです。スイッチのオン・オフで行けるマスが変わるので、状態を持つ必要があります。
そこで、BFSの訪問済みチェックにスイッチの状態も含めて (i, j, sw) の3つで管理します。

from collections import deque

H, W = map(int, input().split())

# 迷路の外枠を壁"#"で囲むことで、境界チェックを楽にする
A = []
A.append(["#"] * (W + 2))
for i in range(H):
    A.append(["#"] + list(input()) + ["#"])
A.append(["#"] * (W + 2))

# 移動パターン(上下左右)
MOVES = [(1, 0), (0, 1), (-1, 0), (0, -1)]

# スタート地点を探す
START = ()
for i in range(H + 2):
    for j in range(W + 2):
        if A[i][j] == "S":
            START = (i, j)
            break

# スタート地点からBFS開始
queue = deque()
queue.append((0, START[0], START[1], False))

visited = set()
visited.add((START[0], START[1], False))

while queue:
    cost, i, j, sw = queue.popleft()

    # ゴールに到達したら、操作回数を出力して終了
    if A[i][j] == "G":
        print(cost)
        exit()

    # 移動パターンに従って、次のマスを探索
    for dx, dy in MOVES:
        ni, nj = i + dx, j + dy

        # 壁または訪問済みならスルー
        if A[ni][nj] == "#" or (ni, nj, sw) in visited:
            continue

        # ゴールに到達したら、操作回数を出力して終了
        if A[ni][nj] == "G":
            print(cost + 1)
            exit()

        # 通常マスまたはスタートマスなら、そのまま移動
        if A[ni][nj] == "." or A[ni][nj] == "S":
            visited.add((ni, nj, sw))
            queue.append((cost + 1, ni, nj, sw))
            continue

        # スイッチマスなら、スイッチの状態を反転させて移動
        if A[ni][nj] == "?":
            visited.add((ni, nj, sw))
            queue.append((cost + 1, ni, nj, not sw))
            continue

        # ドアのマスなら、スイッチの状態に応じて移動可能か判定
        if A[ni][nj] == "o" and sw == False or A[ni][nj] == "x" and sw == True:
            visited.add((ni, nj, sw))
            queue.append((cost + 1, ni, nj, sw))
            continue

# ゴールに到達できなかった場合は-1を出力
print(-1)

E - Reachability Query

E は難しかったので、解説を見て解きました。

UnionFindを使って、連結成分を管理します。

連結成分は共通のリーダーを持つので、black_count[leader] でその連結成分の黒マスの数を管理することができます。

このあたりのUnionFindの仕組みがわからなかったので、問題が解けず...。
単にクラスを使うだけでなく、実装も知っておくほうが良いですね。

# UnionFindは https://qiita.com/sano192/items/027d672129ac4cb42aa2#e---graph-destruction 
# の実装を使わせていただきました。

class UnionFind:
    def __init__(self,n):
        self.n=n
        self.parent_size=[-1]*n
    def leader(self,a):
        if self.parent_size[a]<0: return a
        self.parent_size[a]=self.leader(self.parent_size[a])
        return self.parent_size[a]
    def merge(self,a,b):
        x,y=self.leader(a),self.leader(b)
        if x == y: return
        if abs(self.parent_size[x])<abs(self.parent_size[y]):x,y=y,x
        self.parent_size[x] += self.parent_size[y]
        self.parent_size[y]=x
    def same(self,a,b):
        return self.leader(a) == self.leader(b)
    def size(self,a):
        return abs(self.parent_size[self.leader(a)])
    def groups(self):
        result=[[] for _ in range(self.n)]
        for i in range(self.n):
            result[self.leader(i)].append(i)
        return [r for r in result if r != []]

N, Q = map(int, input().split())

unionfind = UnionFind(N + 1)

# 黒頂点の集合と、各連結成分の黒頂点の数を管理
black_set = set()
black_count = [0] * (N + 1)

for i in range(Q):
    query = list(map(int, input().split()))
    if query[0] == 1:
        # 1のときは、uとvを結合
        u, v = query[1], query[2]

        # すでに連結ならスルー
        if unionfind.same(u, v):
            continue

        # 結合前に、各連結成分の黒頂点の数を取得しておく
        black_count_u = black_count[unionfind.leader(u)]
        black_count_v = black_count[unionfind.leader(v)]

        # uとvを結合
        unionfind.merge(u, v)

        # 新しい連結成分の黒頂点の数を更新
        new_leader = unionfind.leader(u)
        black_count[new_leader] = black_count_u + black_count_v

    if query[0] == 2:
        # 2のときは、頂点vの色を反転
        v = query[1]
        if v in black_set:
            black_set.remove(v)
            black_count[unionfind.leader(v)] -= 1
        else:
            black_set.add(v)
            black_count[unionfind.leader(v)] += 1

    if query[0] == 3:
        # 3のときは、頂点vの連結成分に黒頂点があるかどうかを判定
        v = query[1]
        black_count_v = black_count[unionfind.leader(v)]
        if black_count_v > 0:
            print("Yes")
        else:
            print("No")