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


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