Lilliput Steps

小さな一歩から着実に. 数学やプログラミングのことを書きます.

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);
}