CF375D Tree and Queries -树上启发式合并

思路:
因为是对每个子树都进行统计,所以尝试套用dsu on tree。
对于此题,每个询问还带了一个k是变化的。所以对其进行离线。
离线时很容易想到,我们要保留重链上的每种颜色出现几次的信息col[],但如果是这样还不足以求答案。因为要统计几种颜色符合条件还要遍历所有颜色。线段树可能可以优化,但是有另一种方法。再维护另一个数组res[col[i]]用来保存颜色数大于col[i]的数量,因为col[i]是逐渐递增或递减的,所以可以进行维护,不用担心res中间出现空缺。

//
// Created by acer on 2021/7/6.
//

#include "bits/stdc++.h"

using namespace std;
const int maxn = 3e5;
int cnt = 0;
int head[maxn];
struct node {
    
    
    int to, next;
} a[maxn];

void add(int u, int v) {
    
    
    a[++cnt].to = v;
    a[cnt].next = head[u];
    head[u] = cnt;
}

struct Query {
    
    
    int id;
    int col;
};
vector<Query> query[maxn];
int col[maxn];
int siz[maxn];
int son[maxn];
int mor[maxn];

void dfsFir(int u, int fa) {
    
    
    siz[u] = 1;
    for (int i = head[u]; i; i = a[i].next) {
    
    
        int v = a[i].to;
        if (v == fa) continue;
        dfsFir(v, u);
        siz[u] += siz[v];
        if (siz[v] > siz[son[u]]) {
    
    
            son[u] = v;
        }
    }
}

int flag;
int cntcol[maxn];
int ans[maxn];
int res[maxn];

void count(int u, int fa, int val) {
    
    
    if (val == 1) {
    
    
        cntcol[col[u]] += val;
        res[cntcol[col[u]]] += val;
    } else {
    
    
        res[cntcol[col[u]]] += val;
        cntcol[col[u]] += val;
    }
    for (int i = head[u]; i; i = a[i].next) {
    
    
        int v = a[i].to;
        if (v == fa || v == flag) continue;
        count(v, u, val);
    }
}

void dfsSec(int u, int fa, int keep) {
    
    
    for (int i = head[u]; i; i = a[i].next) {
    
    
        int v = a[i].to;
        if (v == fa || v == son[u]) continue;
        dfsSec(v, u, false);
    }
    if (son[u]) {
    
    
        dfsSec(son[u], u, true);
        flag = son[u];
    }
    count(u, fa, 1);


    for (Query q : query[u]) {
    
    
        ans[q.id] = res[q.col];
    }
    flag = 0;

    if (!keep) {
    
    
        count(u, fa, -1);
    }
}

int main() {
    
    
    ios::sync_with_stdio(0);
//    freopen("in.txt","r",stdin);
    int n, m;
    cin >> n >> m;
    for (int i = 1; i <= n; ++i) {
    
    
        cin >> col[i];
    }
    for (int i = 1; i <= n - 1; ++i) {
    
    
        int u, v;
        cin >> u >> v;
        add(u, v);
        add(v, u);
    }

    for (int i = 1; i <= m; ++i) {
    
    
        int u, c;
        cin >> u >> c;
        query[u].push_back(Query{
    
    i, c});
    }
    dfsFir(1, 0);
    dfsSec(1, 0, 0);
    for (int i = 1; i <= m; ++i) {
    
    
        cout << ans[i] << endl;
    }
}

猜你喜欢

转载自blog.csdn.net/weixin_45509601/article/details/118521976
今日推荐