使用BIT查询彩色树的子树中不同颜色的数量

在树结构中，查询子树中不同颜色的数量是一种常见的问题，可以通过使用基于树状数组（Binary Indexed Tree，BIT）的方法进行优化。

步骤

1. 用哈希表记录每一种颜色的数量，可以将颜色映射为整数。同时，将每个节点的颜色也存储下来。

2. 对于每个节点，求出其子树范围内每种颜色的数量。这可以通过使用BIT进行优化。建立一个数组，记录每个颜色目前为止的数量。然后，遍历该节点的所有子节点，将子节点的颜色数量加到相应的颜色上。

   void dfs(int u) {
       colorCount[color[u]]++;
       for (int i = head[u]; i != -1; i = e[i].next) {
           int v = e[i].to;
           dfs(v);
       }
       for (auto xi: queries[u]) {
           ans[xi.second] = query(xi.first) - ans[xi.second];
       }
   }
   
   void update(int x, int v) {
       for (int i = x; i < N; i += lowbit(i)) {
           colorSum[i] += v;
       }
   }
   
   int query(int x) {
       int res = 0;
       for (int i = x; i > 0; i -= lowbit(i)) {
           res += colorSum[i];
       }
       return res;
   }
   

3. 对于每个查询，先求出其所在节点的颜色数量，再减去其父节点的颜色数量，得到子树中的不同颜色数量。

   for (auto xi: queries[u]) {
       ans[xi.second] = query(xi.first) - ans[xi.second];
   }
   

代码片段

const int N = 100010;
const int M = 400010;

int n, m;
int color[N];
int num[N];
int l[N], r[N];
int cnt = 0;
int lis[2 * N];
int ans[M];

vector<int> g[N];
vector<pair<int, int>> queries[N];

int lowbit(int x) {
    return x & (-x);
}

int colorCount[N], colorSum[N];

void dfs(int u) {
    colorCount[color[u]]++;
    update(color[u], 1);
    l[u] = ++cnt; // 保存起始位置
    for (int i = 0; i < (int)g[u].size(); i++) {
        int v = g[u][i];
        dfs(v);
    }
    r[u] = cnt; // 保存结束位置
    for (auto xi: queries[u]) {
        ans[xi.second] = query(xi.first) - ans[xi.second];
    }
    update(color[u], -1); // 回溯时撤销
}

void update(int x, int v) {
    for (int i = x; i < N; i += lowbit(i)) {
        colorSum[i] += v;
    }
}

int query(int x) {
    int res = 0;
    for (int i = x; i > 0; i -= lowbit(i)) {
        res += colorSum[i];
    }
    return res;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    memset(colorCount, 0, sizeof(colorCount));
    memset(colorSum, 0, sizeof(colorSum));
    cin >> n >> m;
    map<int, int> mp;
    for (int i = 1; i <= n; i++) {
        cin >> color[i];
        if (!mp.count(color[i])) {
            mp[color[i]] = mp.size() + 1; // 离散化
        }
        color[i] = mp[color[i]];
    }
    for (int i = 2; i <= n; i++) {
        int p;
        cin >> p;
        g[p].push_back(i);
    }
    for (int i = 1; i <= m; i++) {
        int u, c;
        cin >> u >> c;
        if (!mp.count(c)) {
            mp[c] = mp.size() + 1; // 离散化
        }
        queries[u].push_back(make_pair(mp[c], i));
    }
    dfs(1);
    for (int i = 1; i <= m; i++) {
        cout << ans[i] << endl;
    }
    return 0;
}