【AtCoder】ABC 448 D - Integer-duplicated Path

2026-03-16

AtCoder/ABC/D問題 AtCoder/ABC/400点問題 AtCoder/茶Diff AtCoder/NoviSteps/1Q AtCoder/グラフ問題/木 AtCoder/DFS 競技プログラミング

D - Integer-duplicated Path

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

AtCoder

実行時間制限: 2 sec / メモリ制限: 1024 MiB / Difficulty: 638 / NoviSteps: 1Q / 配点: 400 点

問題概要

頂点 $1,2,\dots,N$ の $N$ 頂点からなる木が与えられ、この木の辺のうち $i$ 本目は頂点 $U_i$ と頂点 $V_i$ を結んでいる。頂点 $i$ には整数 $A_i$ が書かれている。

全ての $k=1,2,\dots,N$ について以下の問題に答えよ。

頂点 $1$ $1$ から頂点 $k$ $k$ への単純パスに含まれる頂点について、同じ整数の書かれた異なる $2$ $2$ 頂点の組が存在すればYes、そうでないならNoと答えよ。
- なお、木上の $2$ つの頂点を結ぶ単純なパスが一意に定まることは証明できる。

制約

入力は全て整数
$2 \le N \le 2 \times 10^5$
$1 \le A_i \le 10^9$

考察

与えられるグラフは木なので、頂点 $1$ を根とすれば任意の頂点 $k$ までの単純パスは一意に定まり、DFS を用いて全ての頂点に対する根からのパスを計算することができる。

DFS で根からある頂点 $k$ に到達したとき、ここまでのパス上に同じ整数の書かれた異なる $2$ 頂点の組が存在するかどうかを判定するのだが、愚直に行うと $O(N^2)$ の計算量となってしまう。

これを高速化するために、「現在のパス上にどの整数が何個あるか」を記録する連想配列と、「重複している整数の種類数」を記録するカウンターを保持することを考える。その上で、以下のような操作を行う。

ある頂点に到達したとき(行きがけ)
- その頂点の整数を、連想配列に追加する。もし追加する前にすでにその整数の個数が $1$ 個あればカウンターを $1$ 増やす。
ある頂点からその親の頂点に戻るとき(帰りがけ)
- その頂点の整数を、連想配列から減らす。もし減らす前にその整数の個数が $2$ 個あればカウンターを $1$ 減らす。

これにより、頂点に到達した時点でカウンターが $0$ より大きいかどうかをチェックするだけで、判定処理を $O(1)$ の計算量で行うことができる。

したがって、全体としては DFS を行うだけとなり、計算量は $O(N)$ である。

実装例

1.#include <bits/stdc++.h>
2.using namespace std;
3. 
4.#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)
5. 
6.template <class T>
7.using Graph = vector<vector<T>>;
8. 
9.// ======================================== //
10. 
11.int main()
12.{
13.    int N;
14.    cin >> N;
15.    vector<int> A(N);
16.    rep(i, 0, N) cin >> A[i];
17.    Graph<int> G(N);
18.    rep(i, 0, N - 1)
19.    {
20.        int U, V;
21.        cin >> U >> V;
22.        U--, V--;
23.        G[U].push_back(V);
24.        G[V].push_back(U);
25.    }
26. 
27.    map<int, int> mp;
28.    vector<bool> ans(N, false);
29.    int cnt = 0;
30. 
31.    auto dfs = [&](auto self, int now, int parent) -> void
32.    {
33.        if (mp[A[now]] == 1)
34.            cnt++;
35.        mp[A[now]]++;
36. 
37.        if (cnt > 0)
38.            ans[now] = true;
39.        else
40.            ans[now] = false;
41. 
42.        for (auto &&next : G[now])
43.        {
44.            if (next == parent)
45.                continue;
46.            self(self, next, now);
47.        }
48. 
49.        mp[A[now]]--;
50.        if (mp[A[now]] == 1)
51.            cnt--;
52.    };
53. 
54.    dfs(dfs, 0, -1);
55. 
56.    rep(i, 0, N)
57.    {
58.        if (ans[i])
59.            cout << "Yes" << endl;
60.        else
61.            cout << "No" << endl;
62.    }
63. 
64.    return 0;
65.}