[Codeforces] 375D - Tree and Queries

Author            : Dipu Kumar Mohanto 
                    CSE, Batch - 6
                    BRUR.
Problem Statement : 375D - Tree and Queries
Source            : Codeforces
Category          : Graph Theory
Algorithm         : DSU on Tree
Verdict           : Accepted 

#include "bits/stdc++.h"
 
using namespace std;
 
static const int maxn = 1e5 + 5;
 
struct info
{
    int k, id;
    info(int k = 0, int id = 0) : k(k), id(id) {}
};
 
vector <int> graph[maxn];
vector <info> query[maxn];
int color[maxn];
int cnt[maxn], fre[maxn], ans[maxn];  // frequency of color in subtree of v
int subtreeSize[maxn], big[maxn];
 
void dfs(int u = 1, int p = -1)
{
    subtreeSize[u] = 1;
    for (int v : graph[u])
    {
        if (v == p) continue;
        dfs(v, u);
        subtreeSize[u] += subtreeSize[v];
    }
}
 
void add(int u, int p)
{
    cnt[ color[u] ]++;
    fre[ cnt[color[u] ] ]++;
    for (int v : graph[u])
    {
        if (v == p || big[v]) continue;
        add(v, u);
    }
}
 
void remov(int u, int p)
{
    fre[ cnt[ color[u] ] ]--;
    cnt[ color[u] ]--;
    for (int v : graph[u])
    {
        if (v == p || big[v]) continue;
        remov(v, u);
    }
}
 
void dsu(int u = 1, int p = -1, bool keep = 0)
{
    int mx = -1, bigChild = -1;
    for (auto v : graph[u])
    {
        if (v == p) continue;
        if (subtreeSize[v] > mx) mx = subtreeSize[v], bigChild = v;
    }
    for (int v : graph[u])
    {
        if (v == p || v == bigChild) continue;
        dsu(v, u, 0);
    }
    if (bigChild != -1) dsu(bigChild, u, 1), big[bigChild] = 1;
    add(u, p);
    // --- answer the query of subtree u
    for (info it : query[u]) ans[it.id] = fre[it.k];
    if (bigChild != -1) big[bigChild] = 0;
    if (keep == 0) remov(u, p);
}
 
int main()
{
    int tnode, tquery;
    scanf("%d %d", &tnode, &tquery);
    for (int i = 1; i <= tnode; i++) scanf("%d", &color[i]);
    for (int e = 1; e < tnode; e++)
    {
        int u, v;
        scanf("%d %d", &u, &v);
        graph[u].emplace_back(v);
        graph[v].emplace_back(u);
    }
    for (int i = 1; i <= tquery; i++)
    {
        int v, k;
        scanf("%d %d", &v, &k);
        query[v].emplace_back(k, i);
    }
    dfs();
    dsu();
    for (int i = 1; i <= tquery; i++) printf("%d\n", ans[i]);
}