aacord’s memo

abcを中心にpythonで解いた問題のメモ、整理をしています。緑になった。

全国統一プログラミング王決定戦予選 D - Restore the Tree (python)

まず、根になる頂点は、Bには含まれておらず、そのような点 a がただ一つ存在することが分かる。
そして、次に Ai = a となっている辺を全て取り除くと、Bに含まれない点 b が新たに1個以上現れる。その点は、a しか親になりえないので b の親は a とわかる。
つまり、まず根を決めてから、根から順に辺を消す→そこで初めて B に出てくる回数が 0 になった点を見つける → その点を bfs 順にみていき、また辺を消す...を繰り返したらよい。
最初 B の要素をへらすのにセグ木の区間更新を使ったらめちゃしんどくてしかも TLE した

import sys
read = sys.stdin.buffer.read
readline = sys.stdin.buffer.readline
readlines = sys.stdin.buffer.readlines
from collections import deque,defaultdict

n,m = map(int,readline().split())
c = [set() for i in range(n)]
d = defaultdict(int)
for i in range(m+n-1):
  a,b = map(int,readline().split())
  a -= 1
  b -= 1
  c[a].add(b)
  d[b] += 1
  
ans = [0]*n
root = deque()
for i in range(n):
  if d[i] == 0:
    root.append(i)
    break
cnt = 1

while True:
  now = root.popleft()
  for nex in c[now]:
    d[nex] -= 1
    if d[nex] == 0:
      cnt += 1
      ans[nex] = now+1
      root.append(nex)
  if cnt == n:
    break
print(*ans,sep='\n')