はまやんはまやんはまやん

hamayanhamayan's blog

Unique Color [AtCoder Beginner Contest 198 E]

https://atcoder.jp/contests/abc198/tasks/abc198_e

前提知識

  • dfs

解説

https://atcoder.jp/contests/abc198/submissions/21691565

競技プログラミングの問題で木を題材とした問題が多く存在する。
どのようにプログラミングするとよいかは人それぞれかと思うが、隣接グラフを使って頑張ることを勧める。

dfs

今回の問題は木構造のとある根(頂点1)からのある条件を使っている。
よって、頂点1を始点としてdfsすることでうまく扱えるような感じに見える。

cnt[color] := 色colorがこれまでの遷移に何個あるか

これを更新しながらdfsする。
頂点を潜るときにcntの自分の色をインクリメントして、離れるときにデクリメントすることで、この配列に入っている情報を
現在見ている頂点から始点(頂点1)までの情報に限定することができる。
このようなテクは木構造+dfsでよくみられるので慣れると応用範囲が広い。
後はこれを見ながら、同色が他に無ければよい頂点として印をつけていけばいい。

int N, C[101010];
vector<int> E[101010];
//---------------------------------------------------------------------------------------------------
int cnt[101010];
bool good[101010];
void dfs(int cu, int pa = -1) {
    if (cnt[C[cu]] == 0) good[cu] = true;
    cnt[C[cu]]++;

    fore(to, E[cu]) if (to != pa) {
        dfs(to, cu);
    }

    cnt[C[cu]]--;
}
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N;
    rep(i, 0, N) cin >> C[i];
    rep(i, 0, N - 1) {
        int a, b; cin >> a >> b;
        a--; b--;
        E[a].push_back(b);
        E[b].push_back(a);
    }

    dfs(0);
    rep(i, 0, N) if(good[i]) printf("%d\n", i + 1);
}