# Querying the number of distinct colors in a subtree of a colored tree using BIT in C++

In this tutorial, we will be discussing a program to find querying the number of distinct colors in a subtree of a colored tree using BIT.

For this we will be provided with rooted tree where each node has a color denoted by given array. Our task is to find all the distinct coloured nodes below the given node in the tree.

## Example

Live Demo

#include<bits/stdc++.h>
#define MAXIMUM_COLOUR 1000005
#define MAXIMUM_NUMBER 100005
using namespace std;
vector<int> tree[MAXIMUM_NUMBER];
vector<int> table[MAXIMUM_COLOUR];
int isTraversing[MAXIMUM_COLOUR];
int bit[MAXIMUM_NUMBER], getVisTime[MAXIMUM_NUMBER],
getEndTime[MAXIMUM_NUMBER];
int getFlatTree[2 * MAXIMUM_NUMBER];
bool vis[MAXIMUM_NUMBER];
int tim = 0;
vector< pair< pair<int, int>, int> > queries;
//storing results of each queryingTree
int ans[MAXIMUM_NUMBER];
void update(int idx, int val) {
while ( idx < MAXIMUM_NUMBER ) {
bit[idx] += val;
idx += idx & -idx;
}
}
int queryingTree(int idx) {
int result = 0;
while ( idx > 0 ) {
result += bit[idx];
idx -= idx & -idx;
}
return result;
}
void preformingDFS(int v, int color[]) {
//marking the node visited
vis[v] = 1;
getVisTime[v] = ++tim;
getFlatTree[tim] = color[v];
vector<int>::iterator it;
for (it=tree[v].begin(); it!=tree[v].end(); it++)
if (!vis[*it])
preformingDFS(*it, color);
getEndTime[v] = ++tim;
getFlatTree[tim] = color[v];
}
//adding an edge to the tree
void addingNewEdge(int u, int v) {
tree[u].push_back(v);
tree[v].push_back(u);
}
void markingFirstFind(int n) {
for (int i = 1 ; i <= 2 * n ; i++) {
table[getFlatTree[i]].push_back(i);
if (table[getFlatTree[i]].size() == 1) {
update(i, 1);
isTraversing[getFlatTree[i]]++;
}
}
}
void calcQuery() {
int j = 1;
for (int i=0; i<queries.size(); i++) {
for ( ; j < queries[i].first.first ; j++ ) {
int elem = getFlatTree[j];
update( table[elem][isTraversing[elem] - 1], -1);
if ( isTraversing[elem] < table[elem].size() ){
update(table[elem][ isTraversing[elem] ], 1);
isTraversing[elem]++;
}
}
ans[queries[i].second] = queryingTree(queries[i].first.second);
}
}
//counting distinct color nodes
void calcAllColours(int color[], int n, int qVer[], int qn) {
preformingDFS(1, color);
for (int i=0; i<qn; i++)
queries.push_back(make_pair(make_pair(getVisTime[qVer[i]] , getEndTime[qVer[i]]), i) );
sort(queries.begin(), queries.end());
markingFirstFind(n);
calcQuery();
for (int i=0; i<queries.size() ; i++) {
cout << "All distinct colours in the given tree: " << ans[i] << endl;
}
}
int main() {
int number = 6;
int color[] = {0, 2, 3, 3, 4, 1};
int queryVertices[] = {3, 2};
int qn = sizeof(queryVertices)/sizeof(queryVertices[0]);
calcAllColours(color, number, queryVertices, qn);
return 0;
}

## Output

All distinct colours in the given tree: 1
All distinct colours in the given tree: 2

Updated on: 19-Aug-2020

125 Views