HackerRank 101 Hack February - Coloring Tree
問題 : Coloring Tree
概要 :
$n$ 頂点からなる根付き木のそれぞれのノードに色がついている. 頂点$v$ を根とする部分木にある, 異なる色の個数を数えるクエリに$m $ 個答えよ.
$1 \leqq n \leqq 10^5$
$0 \leqq m \leqq 10^5$
$1 \leqq v \leqq n$
$1 \leqq$ (頂点の色) $\leqq 10^9$
解法 :
部分木の色をset で持ちながら高速に併合していきたい.
適当ににset たちをマージすると, 偏ってマージしたときに$O(n^2 \log n)$ 時間でマージされてしまう可能性があり非常に危ない.
そこでデータ構造をマージする一般的なテク(サイズが小さいset を大きいset にマージしていく) を使うことで, この併合作業が$O(n \log ^2 n)$ になる.
前計算すればクエリには$O(1)$ で答えられるから, 全体の計算量は$O(n \log ^2 n + m)$ となる.
コード :
#include <cstdio> #include <vector> #include <set> #include <algorithm> using namespace std; vector<int> G[100000]; int col[100000]; int ans[100000]; set<int> s[100000]; set<int> *dfs(int v, int par = -1) { set<int> *p = &s[v]; p->insert(col[v]); for (int i = 0; i < G[v].size(); i++){ if (G[v][i] != par){ set<int> *q = dfs(G[v][i], v); if (p->size() < q->size()) swap(p, q); while (q->size()){ p->insert(*(q->begin())); q->erase(q->begin()); } } } ans[v] = p->size(); return (p); } int main() { int n, m, r; scanf("%d %d %d", &n, &m, &r); --r; for (int i = 0; i < n - 1; i++){ int a, b; scanf("%d %d", &a, &b); --a; --b; G[a].push_back(b); G[b].push_back(a); } for (int i = 0; i < n; i++){ scanf("%d", col + i); } dfs(r); for (int i = 0; i < m; i++){ int q; scanf("%d", &q); --q; printf("%d\n", ans[q]); } return (0); }