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

コメントを残す

メールアドレスが公開されることはありません。 が付いている欄は必須項目です