HackerEarth El Nino ! problem solution

In this HackerEarth El Nino ! problem solution You have been given a tree T consisting of N nodes rooted at node 1 and an array A consisting of M distinct integers.

Now, you need to find the number of distinct triplets (U, V, K) in tree T such that node V lies in the subtree rooted at node U, and the distance between them is exactly A[K]. We call the distance between 2 nodes to be the number of edges lying in the shortest path among them.

2 triplets (Ui, Vi, Ki) and (Uj, Vj, Kj) are considered to be distinct, if Ui != Uj or Vi !=V j or Ki != Kj.

HackerEarth El Nino ! problem solution.

`import java.io.*;import java.util.*;import java.math.*;public final class solution{    static BufferedReader br=new BufferedReader(new InputStreamReader(System.in));    static FastScanner sc=new FastScanner(br);    static PrintWriter out=new PrintWriter(System.out);    static Random rnd=new Random();    static int[] level,a;    static ArrayList<Integer>[] al;    static int n,m;    static long res=0;        static void shuffle(int[] a)    {        for(int i=0;i<m;i++)        {            int next=rnd.nextInt(m);                        int temp=a[next]; a[next]=a[i]; a[i]=temp;        }    }            static int searchLast(int num)    {        int low=0,high=m-1;            while(low<high)        {            int mid=(low+high+1)>>1;                        if(a[mid]<=num)            {                low=mid;            }            else            {                high=mid-1;            }        }                return (a[low]<=num ? low+1 : 0 );    }        static void dfs(int u,int p,int curr)    {        level[u]=curr;res=res+searchLast(curr-1);                for(int x:al[u])        {            if(x!=p)            {                dfs(x,u,curr+1);            }        }    }        @SuppressWarnings("unchecked")    public static void run() throws Exception    {        n=sc.nextInt();m=sc.nextInt();a=new int[m];                for(int i=0;i<m;i++)        {            a[i]=sc.nextInt();        }                shuffle(a);Arrays.sort(a);al=new ArrayList[n];                for(int i=0;i<n;i++)        {            al[i]=new ArrayList<Integer>();        }                for(int i=1;i<n;i++)        {            int u=sc.nextInt()-1,v=sc.nextInt()-1;                        al[u].add(v);al[v].add(u);        }                level=new int[n];dfs(0,-1,1);out.println(res); out.close();    }        public static void main(String[] args) throws Exception    {        new Thread(null,new Runnable(){            public void run()            {                try                {                    new solution().run();                }                catch(Exception e)                {                                    }            }        },"1",1<<28).start();       }}class FastScanner{    BufferedReader in;    StringTokenizer st;    public FastScanner(BufferedReader in) {        this.in = in;    }        public String nextToken() throws Exception {        while (st == null || !st.hasMoreTokens()) {            st = new StringTokenizer(in.readLine());        }        return st.nextToken();    }        public String next() throws Exception {        return nextToken().toString();    }        public int nextInt() throws Exception {        return Integer.parseInt(nextToken());    }    public long nextLong() throws Exception {        return Long.parseLong(nextToken());    }    public double nextDouble() throws Exception {        return Double.parseDouble(nextToken());    }}`

Second solution

`#include <bits/stdc++.h>using namespace std;const int MAXN = 200005;vector <int> G[MAXN];int A[MAXN];long long int ans = 0;void dfs(int pos, int depth){    ans+=A[depth];    for (int i = 0; i < G[pos].size(); ++i)    {        dfs(G[pos][i],depth+1);    }}int main(){    int n,m;    cin>>n>>m;    for (int i = 0; i < m; ++i)    {        int x;        cin>>x;        if(x <= n)            A[x]++;    }    for (int i = 2; i <= n; ++i)    {        int x;        cin>>x;        G[x].push_back(i);    }    for (int i = 1; i <= n; ++i)    {        A[i]+=A[i-1];    }    dfs(1,0);    cout<<ans<<"\n";    return 0;}`