树上子树颜色询问——C++14版本

主要是感觉 arpa 大佬的例程使用了C++14的特性遍历vector容器,感觉这样写代码更清晰,所以也尝试使用了这种写法。

C++14版本:

#include <bits/stdc++.h>
using namespace std;
const int maxn = 1000010;
int n, m, sz[maxn], c[maxn], cnt[maxn];
bool big[maxn];
vector<int> g[maxn], qid[maxn];
struct Query {
    int c, ans;
} query[maxn];
void add(int u, int p, int x) {
    cnt[ c[u] ] += x;
    for (auto v: g[u])
        if (v != p && !big[v])
            add(v, u, x);
}
void dfs(int u, int p, bool keep) {
    int mx = -1, bigSon = -1;   // mx表示重儿子的sz, bigSon表示重儿子编号
    for (auto v: g[u])
        if (v != p && sz[v] > mx)
            mx = sz[ bigSon = v ];
    for (auto v: g[u])
        if (v != p && v != bigSon)
            dfs(v, u, false);
    if (bigSon != -1)
        dfs(bigSon, u, true),
        big[bigSon] = true;
    add(u, p, 1);
    for (auto id: qid[u])
        query[id].ans = cnt[ query[id].c ];
    if (bigSon != -1)
        big[bigSon] = 0;
    if (!keep)
        add(u, p, -1);
}
int main() {
    cin >> n;
    for (int i = 1; i <= n; i ++) cin >> c[i];
    for (int i = 1; i < n; i ++) {
        int a, b;
        cin >> a >> b;
        g[a].push_back(b);
        g[b].push_back(a);
    }
    cin >> m;
    for (int i = 0; i < m; i ++) {
        int u;
        cin >> u >> query[i].c;
        qid[u].push_back(i);
    }
    dfs(1, -1, false);
    for (int i = 0; i < m; i ++)
        cout << query[i].ans << endl;
    return 0;
}
posted @ 2020-11-04 10:09  quanjun  阅读(116)  评论(0编辑  收藏  举报