Algorithm


Problem Name: Data Structures - Coloring Tree

Problem Link: https://www.hackerrank.com/challenges/coloring-tree/problem?isFullScreen=true

In this HackerRank in Data Structures - Coloring Tree solutions,

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?

Input Format
The first line contains three space separated integers representing the number of nodes in the tree (N), number of queries to answer (M) and the root of the tree.

In each of the next N-1 lines, there are two space separated integers(a b) representing an edge from node a to Node b and vice-versa.

N lines follow: N+ith line contains the color of the ith node.

M lines follow: Each line containg a single integer s.

Output Format
Output exactly M lines, each line containing the output of the ith query.

Constraints
0 <= M <= 105
1 <= N <= 105
1 <= root <= N
1 <= color of the Node <= 109

Example

Sample Input

4 2 1
1 2
2 4
2 3
10
20
20
30
1
2

Sample Output

3
2

Explanation

Query 1-Subtree rooted at 1 is the entire tree and colors used are 10 20 20 30 , so the answer is 3(10,20 and 30)

Query 2-Subtree rooted at 2 contains color 20 20 30, so the answer is 2(20 and 30)

 

 

Code Examples

#1 Code Example with C Programming

Code - 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("%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;
  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;
}
Copy The Code & Try With Live Editor

#2 Code Example with C++ Programming

Code - 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 pll; typedef vector < pair > 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;
}
Copy The Code & Try With Live Editor

#3 Code Example with Java Programming

Code - 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() {
            
            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);
        }

    }
}
Copy The Code & Try With Live Editor

#4 Code Example with Javascript Programming

Code - 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();
}
Copy The Code & Try With Live Editor

#5 Code Example with Python Programming

Code - 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>
Copy The Code & Try With Live Editor
Advertisements

Demonstration


Previous
[Solved] Roy and alpha-beta trees solution in Hackerrank - Hacerrank solution C, C++, java,js, Python
Next
[Solved] Recalling Early Days GP with Trees solution in Hackerrank - Hacerrank solution C, C++, java,js, Python