Edit

【pythonでABC198を解説】E - Unique Color

問題概要

問題ページ

問題文

\(N\) 頂点からなる木が与えられます。\(i\) 番目の辺は頂点 \(A_i\) と頂点 \(B_i\) をつないでいます。頂点 \(i\) は色 \(C_i\) で塗られています (この問題では、色は整数として表されます)。

頂点 \(1\) から頂点 \(x\) への最短パスに、頂点 \(x\) と同じ色で塗られた頂点が頂点 \(x\) 以外に存在しないとき、頂点 \(x\) は よい頂点 であるといいます。

よい頂点を全て求めてください。

制約

  • \(2 \leq N \leq 10^5\)
  • \(1 \leq C_i \leq 10^5\)
  • \(1 \leq A_i,B_i \leq N\)
  • 与えられるグラフは木である
  • 入力は全て整数

問題の考察

ACコード

import sys
import collections


def solve():
    sys.setrecursionlimit(10 ** 6)
    input = sys.stdin.readline
    mod = 10 ** 9 + 7
    n = int(input().rstrip('\n'))
    c = list(map(int, input().rstrip('\n').split()))
    edges = [[] for _ in range(n)]
    for i in range(n - 1):
        edge_a, edge_b = list(map(int, input().rstrip('\n').split()))
        edges[edge_a-1].append(edge_b-1)
        edges[edge_b-1].append(edge_a-1)

    ql = [[0, 0]]
    ql = collections.deque(ql)
    fq = [-1] * n
    fc = collections.defaultdict(int)
    fc[c[0]] += 1
    ans = [1]

    def dfs():
        while ql:
            cost, now_edge = ql.pop()
            if fq[now_edge] < 0:
                fq[now_edge] = 0
                for next_edge in edges[now_edge]:
                    ql.append((cost + 1, next_edge))
                    if fc[c[next_edge]] == 0:
                        ans.append(next_edge + 1)
                    fc[c[next_edge]] += 1
                    dfs()
                    fc[c[next_edge]] -= 1

    dfs()
    print(*sorted(ans), sep="\n")


if __name__ == '__main__':
    solve()

-Edit
-,