AtCoder Beginner Contest 216 D – Pair of Balls をPython3で解く

Share

AtCoder上にある問題のうち、AtCoder Problemsでdiff 800以上と判定されているものを順番に解いていく企画。
基本的な考え方は全てコード中のコメントに入れてあるので、参照のこと。

出典:
AtCoder Beginner Contest 216 D – Pair of Balls

愚直にシミュレートすれば解けるが、ペアを効率よく見つけられるように実装する必要がある。

# AtCoder Beginner Contest 216 D - Pair of Balls
# https://atcoder.jp/contests/abc216/tasks/abc216_d
# tag: 愚直 グラフ トポロジカルソート

def main():
    N, M = map(int, input().split())
    tubes = []
    for _ in range(M):
        k = int(input())

        # pop() で処理しやすいように逆順にしておく。
        tubes.append(list(map(int, input().split()))[::-1])

    # 現在ペアになっている色。
    paired = []

    # 色別に取り出せる位置を管理。
    available = [[] for _ in range(N+1)]

    # 初回準備。
    for pos, colors in enumerate(tubes):
        if len(tubes) > 0:
            color = colors[-1]
            available[color].append(pos)
            if len(available[color]) == 2:
                paired.append(color)

    # 取り除いたボールの個数を数えておく。
    removed = 0

    # ここから具体的な操作のシミュレートを行う。
    while len(paired) > 0:
        # ペアに出来る色を取り出し、新たに現れる色を管理する。
        color = paired.pop()
        removed += 1
        for pos in available[color]:
            # 取り出し。
            tubes[pos].pop()
            if len(tubes[pos]) == 0:
                continue

            # まだボールがあるなら、available を更新する。
            new_color = tubes[pos][-1]
            available[new_color].append(pos)

            # 一番上に 2つあるなら、取り除くペアに追加。
            if len(available[new_color]) == 2:
                paired.append(new_color)

    if removed == N:
        print('Yes')
    else:
        print('No')

main()

一方で、グラフの問題に帰着することも出来るのが面白い。

# ある色のボールを取り除ける条件は、上にあるボールが
# 取り除かれていることである。

# そこで、各色を頂点として、各色からひとつ下にある
# 色に対して辺を張るようなグラフを構築する。

# 入辺が存在しない頂点と、そこから出ている辺を順に
# 取り除いていき、全ての頂点を取り除けるならば
# 答えは 'Yes' になる。

# ちなみに、これはトポロジカルソートと同じ操作になっている。

def main():
    N, M = map(int, input().split())
    tubes = []
    for _ in range(M):
        k = int(input())
        tubes.append(list(map(int, input().split())))

    removed = []

    # 頂点ごとの入辺の数。
    edge_in = [0] * (N+1)

    # グラフを構築する。
    paths = [[] for _ in range(N+1)]
    for balls in tubes:
        for i in range(len(balls)-1):
            paths[balls[i]].append(balls[i+1])
            edge_in[balls[i+1]] += 1

    # 入辺の無いものを抽出。
    no_edge_in = [idx for idx, n in enumerate(edge_in) if n == 0]

    # 入辺が無いものと、そこからの辺をどんどん取り除いていく。
    while no_edge_in:
        color = no_edge_in.pop()
        removed.append(color)
        for nxt in paths[color]:
            edge_in[nxt] -= 1
            if edge_in[nxt] == 0:
                no_edge_in.append(nxt)
    
    # 1-indexed にしている関係で、取り除かれる頂点は N+1 個。
    if len(removed) == N + 1:
        print('Yes')
    else:
        print('No')

main()
Share

コメントを残す

メールアドレスが公開されることはありません。