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;
}