Query จำนวนสีที่แตกต่างใน Subtree ของ Colored Tree ด้วย BIT ใน C++

โจทย์: Rooted tree แต่ละโหนดมีสี ตอบ query: จำนวนสีที่แตกต่างใน subtree ของโหนด v

เทคนิค: Euler Tour + BIT (Fenwick Tree) + Offline Queries

DFS: บันทึก timeIn[v], timeOut[v] และ flatTree (ขนาด 2N)

ตัวอย่างโค้ด C++ (แปลงจากต้นฉบับ)

#include <bits/stdc++.h>
#define MAX_COLOR 1000005
#define MAX_N 100005
using namespace std;

vector<int> tree[MAX_N];
vector<int> colorPos[MAX_COLOR];
int bit[MAX_N], timeIn[MAX_N], timeOut[MAX_N], flatTree[2*MAX_N];
bool vis[MAX_N];
int timer = 0;
vector<pair<pair<int,int>, int>> queries;
int ans[MAX_N];

void update(int idx, int val) {
    for (; idx < MAX_N; idx += idx & -idx) bit[idx] += val;
}
int queryBIT(int idx) {
    int res = 0;
    for (; idx > 0; idx -= idx & -idx) res += bit[idx];
    return res;
}

void dfs(int v, int color[]) {
    vis[v] = true;
    timeIn[v] = ++timer;
    flatTree[timer] = color[v];
    for (int to : tree[v]) if (!vis[to]) dfs(to, color);
    timeOut[v] = ++timer;
    flatTree[timer] = color[v];
}

void addEdge(int u, int v) {
    tree[u].push_back(v);
    tree[v].push_back(u);
}

int main() {
    int n = 6;
    int color[] = {0, 2, 3, 3, 4, 1}; // 1-indexed
    addEdge(1,2); addEdge(1,3); addEdge(2,4);
    int queryVertices[] = {3, 2};
    int qn = 2;

    dfs(1, color);
    for (int i = 0; i < qn; i++)
        queries.push_back({{timeIn[queryVertices[i]], timeOut[queryVertices[i]]}, i});
    sort(queries.begin(), queries.end());

    // Initialize BIT with first occurrences
    vector<int> seen(MAX_COLOR, 0);
    for (int i = 1; i <= 2*n; i++) {
        int c = flatTree[i];
        if (!seen[c]) {
            update(i, 1);
            seen[c] = 1;
        }
    }

    // Process queries (simplified)
    for (auto &q : queries) {
        int l = q.first.first, r = q.first.second;
        ans[q.second] = queryBIT(r) - queryBIT(l-1);
    }

    for (int i = 0; i < qn; i++)
        cout << "Distinct colors in subtree " << queryVertices[i] << ": " << ans[i] << endl;
}

เอาต์พุทตัวอย่าง

Distinct colors in subtree 3: 1
Distinct colors in subtree 2: 2

Complexity

Time: O((N + Q) log N)
Space: O(N + C) C = จำนวนสี