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() def AddNode(self, label1, label2): 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: x.add(_) 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) t.AddNode(v1, 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.io.InputStreamReader; 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 { BufferedReader br = new BufferedReader(new InputStreamReader(is, Charset.defaultCharset()), 1 << 17); 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; queries.add(new Query(tree.nodes[qIdx].idx,tree.nodes[qIdx].max)); } List<Query> processOrdered=new ArrayList<>(); processOrdered.addAll(queries); 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--; add(colorArr[l]); } for(;h<query.h;h++){ add(colorArr[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--; } } private void add(int i) { 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<>(); s.add(nodes[i]); 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; s.add(a); } else { throw new RuntimeException("unexp"); } } break; case LAST: e.max = idx; s.pop(); break; default: break; } } } public void addEdge(int a, int b) { nodes[b].addNeigh(nodes[a]); nodes[a].addNeigh(nodes[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<>(); public void addNeigh(Node node) { neigh.add(node); } } }
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("%dn", 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}; tree_node all[100000],*head[100000]; 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]; head[i]=NULL; ans[i]=0; trace[i]=0; tree_size[i]=0; } dfs(root-1); while(M--){ scanf("%d",&x); printf("%dn",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; insert(head+x,all+x); tree_size[x]=1; for(node=table[x];node;node=node->next) if(!trace[node->x]){ dfs(node->x); p1=x; p2=node->x; t1=(tree_size[p1]>tree_size[p2])?&head[p1]:&head[p2]; t2=(tree_size[p1]>tree_size[p2])?head[p2]:head[p1]; big=(tree_size[p1]>tree_size[p2])?p1:p2; tra(t2,t1,big); if(big!=p1){ head[p1]=head[p2]; 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(); }); function readLine() { 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, adjacent: [] }; })); tree.forEach(function(edge) { [0, 1].forEach(function(direction) { nodes[edge[direction]].adjacent.push(nodes[edge[direction ? 0 : 1]]); }); }); 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); parentNode.children = parentNode.adjacent.filter(function(childNode) { return !nodesSeen[childNode.index]; }); delete parentNode.adjacent; 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; curNode.distinctColorsSearch = { 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) { curNode.distinctColorsSearch.colorsToIgnore[colorToIgnore] && curNode.numDistinctColorsInSubtree--; !curNode.distinctColorsSearch.colorsToIgnore[colorToIgnore] && (curNode.distinctColorsSearch.colorsToIgnore[colorToIgnore] = 0); curNode.distinctColorsSearch.colorsToIgnore[colorToIgnore] += childNode.distinctColorsSearch.colorsToIgnore[colorToIgnore]; if (curNode.distinctColorsSearch.colorsToIgnore[colorToIgnore] == colorCounts[colorToIgnore]) { delete curNode.distinctColorsSearch.colorsToIgnore[colorToIgnore]; } } 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 nmr = readLine().split(' '); 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++) { const colorItem = parseInt(readLine(), 10); color.push(colorItem); } let s = []; for (let sItr = 0; sItr < m; sItr++) { const sItem = parseInt(readLine(), 10); s.push(sItem); } let result = solve(tree, color, r, s); ws.write(result.join("n") + "n"); ws.end(); }