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};
   addingNewEdge(1, 2);
   addingNewEdge(1, 3);
   addingNewEdge(2, 4);
   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