# HackerRank Coloring Tree problem solution

In this HackerRank Coloring Tree problem solution, you are given a tree with N nodes with every node being colored. A color is represented by an integer ranging from 1 to 109. Can you find the number of distinct colors available in a subtree rooted at the node s?

## Problem solution in Python programming.

```class Stack:
"A container with a last-in-first-out (LIFO) queuing policy."
def __init__(self):
self.list = []

def push(self,item):
"Push 'item' onto the stack"
self.list.append(item)

def pop(self):
"Pop the most recently pushed item from the stack"
return self.list.pop()

def isEmpty(self):
"Returns true if the stack is empty"
return len(self.list) == 0

def pick(self):
return self.list[len(self.list) - 1]

class Queue:
"A container with a first-in-first-out (FIFO) queuing policy."
def __init__(self):
self.list = []

def push(self,item):
"Enqueue the 'item' into the queue"
self.list.insert(0,item)

def pop(self):
"""
Dequeue the earliest enqueued item still in the queue. This
operation removes the item from the queue.
"""
return self.list.pop()

def isEmpty(self):
"Returns true if the queue is empty"
return len(self.list) == 0

class Tree:

def __init__(self):
self.Nodes = dict()

if label1 not in self.Nodes:
self.Nodes[label1] = Node(label1, set(), list(), None)

if label2 not in self.Nodes:
self.Nodes[label2] = Node(label2, set(), list(), None)

node1 = self.Nodes[label1]
node2 = self.Nodes[label2]

node1.Nodes.append(node2)
node2.Nodes.append(node1)

def SetRoot(self, root):
s = Stack()
s.push(root)

while not s.isEmpty():
node = s.pick()

if(node.Nodes == None):
s.pop()
node.Count = len(node.Tag)
if (node.Parent != None):
node.Parent.Tag = join(node.Parent.Tag, node.Tag)
node.Tag = None
else:
for n in node.Nodes:
n.Nodes.remove(node)
n.Parent = node
s.push(n)
node.Nodes = None

def ColorPathFromNode(self, node, color):
l = 0
while node is not None and node.Visited != color:
node.Visited = color
node.Count = node.Count + 1
node = node.Parent
l = l + 1

def join(x=set(), y=set()):
if len(x) < len(y):
return join(y, x)
for _ in y:
y.clear()
return x

class Node:

def __init__(self, label, tag, nodes, color):
self.Label = label
self.Tag = tag
self.Nodes = nodes
self.Color = color
self.Visited = 0
self.Parent = None
self.Count = 0

n, m, s = input().split()
n, m, s = int(n), int(m), int(s)

t = Tree()

for i in range(0, n - 1):
v1, v2 = input().split()
v1, v2 = int(v1), int(v2)

for i in range(0, n):
color = int(input())
t.Nodes[i + 1].Tag = set([color])

t.SetRoot(t.Nodes[s])

for i in range(0, m):
x = int(input())
print(t.Nodes[x].Count)```

## Problem solution in Java Programming.

```import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.InputStream;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Scanner;
import java.util.Set;
import java.util.Stack;

public class Solution {
public static void main(String args[]) throws Exception {
InputStream is;
if (System.getProperty("user.name").equals("kirk")) {
is = new FileInputStream("input/ds/coloring-tree/input0.txt");
} else {
is = System.in;
}
runCases(is);
// System.out.println(c.result);
}

private static void runCases(InputStream is) throws Exception {
Scanner sc = new Scanner(br);

int N = sc.nextInt();
int Q = sc.nextInt();
int R = sc.nextInt();
RootedTree tree = new RootedTree(N);
for (int i = 0; i < N - 1; i++) {
int a = sc.nextInt();
int b = sc.nextInt();
tree.addEdge(a - 1, b - 1);
}
for (int i = 0; i < N ; i++) {
int c = sc.nextInt() - 1;
tree.setColor(i,c);
}
tree.anchor(R-1);

List<Query    > queries=new ArrayList<>();
for (int i = 0; i < Q; i++) {
int qIdx = sc.nextInt() - 1;
}
List<Query>    processOrdered=new ArrayList<>();
processOrdered.sort(Query.SQORDER);

QueryRunner qr = new QueryRunner(tree.getColorArr());
for (Query query : processOrdered) {
qr.runQuery(query);
}
for (Query query : queries) {
System.out.println(query.result);
}
// System.out.println(asi.result);
}

static class Query{

public static final Comparator<Query> SQORDER = new Comparator<Query>() {

static final int W=3000;
@Override
public int compare(Query o1, Query o2) {
int rl = Integer.compare(o1.l/3000,o2.l/3000);
if(rl!=0){
return rl;
}
return Integer.compare(o1.h, o2.h);
}
};
private int l;
private int h;
private int result;

public Query(int l, int h) {
this.l = l;
this.h = h;
}
};

static enum NodeState {
FIRST, LAST
};

static class QueryRunner{

int cnt[];
int    l,h;
private int[] colorArr;
private int currentCard;

public QueryRunner(int[] colorArr) {
this.colorArr = colorArr;
cnt=new int[colorArr.length];
}

public void runQuery(Query query) {
for(;l>query.l;){
l--;
}
for(;h<query.h;h++){
}
for(;l<query.l;l++){
sub(colorArr[l]);
}
for(;h>query.h;){
h--;
sub(colorArr[h]);
}
query.result=currentCard;

}

private void sub(int i) {
cnt[i]--;
if(cnt[i] == 0){
currentCard--;
}
}

if(cnt[i] == 0){
currentCard++;
}
cnt[i]++;
}

}
static class RootedTree {

private Node[] nodes;
private int nColors;

public RootedTree(int n) {
nodes = new Node[n];
for (int i = 0; i < n; i++) {
nodes[i] = new Node();
}
}

public int[]    getColorArr(){
int    ret[]=new int[nodes.length];
for (Node n : nodes) {
ret[n.idx]=n.color;
}
return ret;
}

Map<Integer,Integer>    cMap=new HashMap<>();
public void setColor(int i, int nextInt) {
Integer cV = cMap.get(nextInt);
if(cV==null){
nodes[i].color=nColors;
cMap.put(nextInt, nColors++);
}else{
nodes[i].color=cV;
}
}

public int query(int i, int j) {
Node cp = commonParent(nodes[i], nodes[j]);
return Math.max(maxV(nodes[i], cp), maxV(nodes[j], cp));
}

private Node commonParent(Node a, Node b) {
if (a.depth > b.depth) {
return commonParent(b, a);
}
while (a.max <= b.idx || b.idx < a.idx) {
a = a.parent;
}
return a;
}

private int maxV(Node node, Node cp) {
int v = 0;
boolean replace = true;
while (node != null) {
v += node.value;
if (replace && node.value > v) {
v = node.value;
}
if (node == cp) {
replace = false;
}
node = node.parent;
}
return v;
}

public void add(int i, int nextInt) {
nodes[i].value += nextInt;
int w = nodes[i].max - nodes[i].idx;
}

public void anchor(int i) {
Stack<Node> s = new Stack<>();

nodes[i].depth = 0;
int idx = 0;
while (!s.isEmpty()) {
Node e = s.peek();
switch (e.state) {
case FIRST:
e.idx = idx++;
e.state = NodeState.LAST;
e.neigh.remove(e.parent);
for (Node a : e.neigh) {
if (a.state == NodeState.FIRST) {
a.depth = e.depth + 1;
a.parent = e;
} else {
throw new RuntimeException("unexp");
}
}
break;
case LAST:
e.max = idx;
s.pop();
break;
default:
break;
}
}
}

public void addEdge(int a, int b) {
}
};

static class Node {
public int color;
public Node parent;
public int value;
public int depth;
public int max;
public int idx;
public NodeState state = NodeState.FIRST;
Set<Node> neigh = new HashSet<>();

}

}
}```

## Problem solution in C++ programming.

```#define _CRT_SECURE_NO_WARNINGS
#include <string>
#include <vector>
#include <algorithm>
#include <numeric>
#include <set>
#include <map>
#include <queue>
#include <iostream>
#include <sstream>
#include <cstdio>
#include <cmath>
#include <ctime>
#include <cstring>
#include <cctype>
#include <cassert>
#define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i))
#define rer(i,l,u) for(int (i)=(int)(l);(i)<=(int)(u);++(i))
#define reu(i,l,u) for(int (i)=(int)(l);(i)<(int)(u);++(i))
#if defined(_MSC_VER) || __cplusplus > 199711L
#define aut(r,v) auto r = (v)
#else
#define aut(r,v) typeof(v) r = (v)
#endif
#define each(it,o) for(aut(it, (o).begin()); it != (o).end(); ++ it)
#define all(o) (o).begin(), (o).end()
#define pb(x) push_back(x)
#define mp(x,y) make_pair((x),(y))
#define mset(m,v) memset(m,v,sizeof(m))
#define INF 0x3f3f3f3f
#define INFL 0x3f3f3f3f3f3f3f3fLL
using namespace std;
typedef vector<int> vi; typedef pair<int,int> pii; typedef vector<pair<int,int> > vpii;
typedef long long ll; typedef vector<long long> vl; typedef pair<long long,long long> pll; typedef vector<pair<long long,long long> > vpll;
typedef vector<string> vs; typedef long double ld;
template<typename T, typename U> inline void amin(T &x, U y) { if(y < x) x = y; }
template<typename T, typename U> inline void amax(T &x, U y) { if(x < y) x = y; }

vector<vi> g;
vector<int> parent;
vi t_ord;

void tree_getorder(int root) {
int n = g.size();
parent.assign(n, -1);

vector<int> stk; stk.push_back(root);
while(!stk.empty()) {
int i = stk.back(); stk.pop_back();
t_ord.push_back(i);
rep(j, g[i].size()) {
int c = g[i][j];
if(parent[c] == -1 && c != root) {
stk.push_back(c);
}else {
parent[i] = c;
}
}
}
}

int ans[100000];
set<int> s[100000];
int color[100000];
int main() {
int N, M, root;
scanf("%d%d%d", &N, &M, &root); root --;
g.assign(N, vi());
rep(i, N-1) {
int a, b;
scanf("%d%d", &a, &b);
a --, b --;
g[a].pb(b);
g[b].pb(a);
}
rep(i, N) scanf("%d", &color[i]);
tree_getorder(root);
rep(i, N) s[i].clear();
for(int ordi = N-1; ordi >= 0; ordi --) {
int v = t_ord[ordi];
s[v].insert(color[v]);
ans[v] = s[v].size();
int u = parent[v];
if(u == -1) continue;
if(s[u].size() < s[v].size())
swap(s[u], s[v]);
s[u].insert(all(s[v]));
set<int>().swap(s[v]);
}
rep(i, M) {
int s;
scanf("%d", &s); s --;
printf("%d\n", ans[s]);
}
return 0;
}```

## Problem solution in C programming.

```#include <stdio.h>
#include <stdlib.h>
typedef struct _list{
int x;
struct _list *next;
} list;
typedef struct _tree_node{
enum {red,black} colour;
int data;
struct _tree_node *left,*right,*parent;
} tree_node;
void sort_a2(int*a,int*b,int size);
void merge2(int*a,int*left_a,int*right_a,int*b,int*left_b,int*right_b,int left_size, int right_size);
void insert_edge(int x,int y);
void dfs(int x);
void tra(tree_node *t2,tree_node **t1,int big);
void left_rotate(tree_node **root,tree_node *x);
void right_rotate(tree_node **root,tree_node *y);
void reconstruct(tree_node **root,tree_node *x);
int normal_insert(tree_node **root,tree_node *x);
int insert(tree_node **root,tree_node *x);
int color[100000],c[100000],ans[100000],cc[100000],idx[100000],trace[100000],tree_size[100000];
list *table[100000]={0};

int main(){
int N,M,root,x,y,i;
scanf("%d%d%d",&N,&M,&root);
for(i=0;i<N-1;i++){
scanf("%d%d",&x,&y);
insert_edge(x-1,y-1);
}
for(i=0;i<N;i++){
scanf("%d",color+i);
idx[i]=i;
cc[i]=0;
}
sort_a2(color,idx,N);
for(i=x=0;i<N;i++){
if(i && color[i]!=color[i-1])
x++;
c[idx[i]]=x;
cc[x]++;
}
for(i=0;i<N;i++){
all[i].data=c[i];
ans[i]=0;
trace[i]=0;
tree_size[i]=0;
}
dfs(root-1);
while(M--){
scanf("%d",&x);
printf("%d\n",ans[x-1]);
}
return 0;
}
void sort_a2(int*a,int*b,int size){
if (size < 2)
return;
int m = (size+1)/2,i;
int*left_a,*left_b,*right_a,*right_b;
left_a=(int*)malloc(m*sizeof(int));
right_a=(int*)malloc((size-m)*sizeof(int));
left_b=(int*)malloc(m*sizeof(int));
right_b=(int*)malloc((size-m)*sizeof(int));
for(i=0;i<m;i++){
left_a[i]=a[i];
left_b[i]=b[i];
}
for(i=0;i<size-m;i++){
right_a[i]=a[i+m];
right_b[i]=b[i+m];
}
sort_a2(left_a,left_b,m);
sort_a2(right_a,right_b,size-m);
merge2(a,left_a,right_a,b,left_b,right_b,m,size-m);
free(left_a);
free(right_a);
free(left_b);
free(right_b);
return;
}
void merge2(int*a,int*left_a,int*right_a,int*b,int*left_b,int*right_b,int left_size, int right_size){
int i = 0, j = 0;
while (i < left_size|| j < right_size) {
if (i == left_size) {
a[i+j] = right_a[j];
b[i+j] = right_b[j];
j++;
} else if (j == right_size) {
a[i+j] = left_a[i];
b[i+j] = left_b[i];
i++;
} else if (left_a[i] <= right_a[j]) {
a[i+j] = left_a[i];
b[i+j] = left_b[i];
i++;
} else {
a[i+j] = right_a[j];
b[i+j] = right_b[j];
j++;
}
}
return;
}
void insert_edge(int x,int y){
list *node;
node=(list*)malloc(sizeof(list));
node->x=x;
node->next=table[y];
table[y]=node;
node=(list*)malloc(sizeof(list));
node->x=y;
node->next=table[x];
table[x]=node;
return;
}
void dfs(int x){
int p1,p2,big;
list *node;
tree_node **t1,*t2;
trace[x]=1;
all[x].left=all[x].right=all[x].parent=NULL;
tree_size[x]=1;
for(node=table[x];node;node=node->next)
if(!trace[node->x]){
dfs(node->x);
p1=x;
p2=node->x;
big=(tree_size[p1]>tree_size[p2])?p1:p2;
tra(t2,t1,big);
if(big!=p1){
tree_size[p1]=tree_size[p2];
}
}
ans[x]=tree_size[x];
return;
}
void tra(tree_node *t2,tree_node **t1,int big){
if(!t2)
return;
tra(t2->left,t1,big);
tra(t2->right,t1,big);
t2->parent=t2->left=t2->right=NULL;
if(insert(t1,t2))
tree_size[big]++;
return;
}
void left_rotate(tree_node **root,tree_node *x){
tree_node *y;
y=x->right;
if(!y) return;
x->right=y->left;
if(y->left)
y->left->parent=x;
y->parent=x->parent;
if(x->parent==NULL) *root=y;
else
if(x==x->parent->left)
x->parent->left=y;
else
x->parent->right=y;
y->left=x;
x->parent=y;
return;
}
void right_rotate(tree_node **root,tree_node *y){
tree_node *x;
x=y->left;
if(!x) return;
y->left=x->right;
if(x->right)
x->right->parent=y;
x->parent=y->parent;
if(y->parent==NULL) *root=x;
else
if(y==y->parent->right)
y->parent->right=x;
else
y->parent->left=x;
x->right=y;
y->parent=x;
return;
}
void reconstruct(tree_node **root,tree_node *x){
tree_node *y,*z;
y=x->parent;
z=x->parent->parent;
x->colour=black;
z->colour=red;
x->parent=z->parent;
if(z->parent==NULL)
*root=x;
else if(z==z->parent->left)
z->parent->left=x;
else
z->parent->right=x;
if(z->left==y){
x->left=y;
x->right=z;
}
else{
x->left=z;
x->right=y;
}
y->parent=z->parent=x;
y->left=y->right=z->left=z->right=NULL;
return;
}
int normal_insert(tree_node **root,tree_node *x){
if(*root==NULL)
*root=x;
else if((*root)->data==x->data)
return 0;
else{
x->parent=*root;
if((*root)->data>x->data)
return normal_insert(&((*root)->left),x);
else
return normal_insert(&((*root)->right),x);
}
return 1;
}
int insert(tree_node **root,tree_node *x){
if(!normal_insert(root,x))
return 0;
tree_node *y;
x->colour=red;
while(x!=*root && x->parent->colour==red){
if(x->parent==x->parent->parent->left){
y=x->parent->parent->right;
if(!y)
if(x==x->parent->left){
x->parent->colour=black;
x->parent->parent->colour=red;
right_rotate(root,x->parent->parent);
}
else{
y=x->parent;
reconstruct(root,x);
x=y;
}
else if(y->colour==red){
x->parent->colour=black;
y->colour=black;
x->parent->parent->colour=red;
x=x->parent->parent;
}
else{
if(x==x->parent->right){
x=x->parent;
left_rotate(root,x);
}
x->parent->colour=black;
x->parent->parent->colour=red;
right_rotate(root,x->parent->parent);
}
}
else{
y=x->parent->parent->left;
if(!y)
if(x==x->parent->right){
x->parent->colour=black;
x->parent->parent->colour=red;
left_rotate(root,x->parent->parent);
}
else{
y=x->parent;
reconstruct(root,x);
x=y;
}
else if(y->colour==red){
x->parent->colour=black;
y->colour=black;
x->parent->parent->colour=red;
x=x->parent->parent;
}
else{
if(x==x->parent->left){
x=x->parent;
right_rotate(root,x);
}
x->parent->colour=black;
x->parent->parent->colour=red;
left_rotate(root,x->parent->parent);
}
}
}
(*root)->colour=black;
return 1;
}```

## Problem solution in JavaScript programming.

```'use strict';

const fs = require('fs');

process.stdin.resume();
process.stdin.setEncoding('utf-8');

let inputString = '';
let currentLine = 0;

process.stdin.on('data', inputStdin => {
inputString += inputStdin;
});

process.stdin.on('end', _ => {
inputString = inputString.trim().split('\n').map(str => str.trim());

main();
});

return inputString[currentLine++];
}

// Complete the solve function below.
function solve(tree, color, r, s) {
var nodes = [null].concat(color.map(function(nodeColor, colorIndex) {
return {
index: colorIndex + 1,
color: nodeColor,
};
}));
tree.forEach(function(edge) {
[0, 1].forEach(function(direction) {
});
});
var rootNode = nodes[r];
var nodesByTreeOrder = [];
function directEdges(rootNode) {
var nodesSeen = {};

var nodesToTraverse = [rootNode];
while (nodesToTraverse.length) {
var parentNode = nodesToTraverse.pop();
if (nodesSeen[parentNode.index]) {
continue;
}
nodesSeen[parentNode.index] = true;

nodesByTreeOrder.push(parentNode);

return !nodesSeen[childNode.index];
});

parentNode.children.forEach(function(childNode) {
childNode.parent = parentNode;

nodesToTraverse.push(childNode);
});
}
};
directEdges(rootNode);
nodesByTreeOrder = nodesByTreeOrder.reverse();

var colorCounts = {};
color.forEach(function(color) {
!colorCounts[color] && (colorCounts[color] = 0);
colorCounts[color]++;
});

nodesByTreeOrder.forEach(function(curNode) {
curNode.numDistinctColorsInSubtree = 1;
colorsToIgnore: {}
};
(colorCounts[curNode.color] > 1) && (curNode.distinctColorsSearch.colorsToIgnore[curNode.color] = 1);

curNode.children.forEach(function(childNode) {
curNode.numDistinctColorsInSubtree += childNode.numDistinctColorsInSubtree;

for (var colorToIgnore in childNode.distinctColorsSearch.colorsToIgnore) {

}
}

delete childNode.distinctColorSearch;
});
});
delete rootNode.distinctColorSearch;

return s.map(function(fromNodeIndex) {
return nodes[fromNodeIndex].numDistinctColorsInSubtree;
});
}

function main() {
const ws = fs.createWriteStream(process.env.OUTPUT_PATH);

const n = parseInt(nmr[0], 10);

const m = parseInt(nmr[1], 10);

const r = parseInt(nmr[2], 10);

let tree = Array(n-1);

for (let treeRowItr = 0; treeRowItr < n-1; treeRowItr++) {
tree[treeRowItr] = readLine().split(' ').map(treeTemp => parseInt(treeTemp, 10));
}

let color = [];

for (let colorItr = 0; colorItr < n; colorItr++) {
color.push(colorItem);
}

let s = [];

for (let sItr = 0; sItr < m; sItr++) {
s.push(sItem);
}

let result = solve(tree, color, r, s);

ws.write(result.join("\n") + "\n");

ws.end();
}```