HackerEarth Velma and Queries problem solution

In this HackerEarth Velma and Queries problem solution The "Mystery Incorporated" once went for an investigation. Since Scooby and Shaggy always create havoc during investigation, Velma decided to give them a question so as to keep them busy.

Velma gave them a tree consisting of N nodes, rooted at node 1 with a number in each node, and asked them several queries. Each query was of the form u l and they had to find the sum of values of all nodes situated at a depth l and were in the subtree of u.

Since Scooby and Shaggy are dumb, they need your help in solving this problem.

HackerEarth Velma and Queries problem solution.

#include <bits/stdc++.h>
using namespace std;

const int N = 1e5;
vector<int > tree[N + 5], nodes_at_level[N + 5];
vector<long long > values_at_level[N + 5];
long long arr[N + 5];
int level[N + 5], tin[N + 5], tout[N + 5], timer;


void dfs(int u, int lev = 1, int p = 1)
{
  tin[u] = timer++;
  level[u] = lev;
  nodes_at_level[lev].push_back(tin[u]);
  values_at_level[lev].push_back(arr[u]);
  for (int i = 0;i < tree[u].size();i++)
  {
    int to = tree[u][i];
    if (to != p)
      dfs(to, lev + 1, u);
  }
  tout[u] = timer++;
}

void compute_prefix()
{
  for (int i = 1;i <= N;i++)
  {
    for (int j = 1;j < values_at_level[i].size();j++)
      values_at_level[i][j] += values_at_level[i][j - 1];
  }
}

int main()
{
  int n, q, u, v, l;
  scanf("%d%d", &n, &q);
  for (int i = 1;i <= n;i++)
    scanf("%lld", &arr[i]);
  for (int i = 0;i < n - 1;i++)
  {
    scanf("%d%d", &u, &v);
    tree[u].push_back(v);
    tree[v].push_back(u);
  }

  dfs(1);
  compute_prefix();

  while (q--)
  {
    scanf("%d%d", &u, &l);
    int lt, rt;
    lt = lower_bound(nodes_at_level[l].begin(), nodes_at_level[l].end(), tin[u])-nodes_at_level[l].begin();
    rt= lower_bound(nodes_at_level[l].begin(), nodes_at_level[l].end(), tout[u]) - nodes_at_level[l].begin();
    lt--;rt--;
    long long ans;
    if (rt < 0)ans = 0;
    else if (lt < 0)ans = values_at_level[l][rt];
    else ans = values_at_level[l][rt] - values_at_level[l][lt];
    printf("%lld\n", ans);
  }
  return 0;
}

Second solution

#include <string>
#include <vector>
#include <map>
#include <list>
#include <iterator>
#include <cassert>
#include <set>
#include <queue>
#include <iostream>
#include <sstream>
#include <stack>
#include <deque>
#include <cmath>
#include <memory.h>
#include <cstdlib>
#include <cstdio>
#include <cctype>
#include <algorithm>
#include <utility>
#include <time.h>
#include <complex>
using namespace std;

#define FOR(i, a, b) for(int i=(a);i<(b);i++)
#define RFOR(i, b, a) for(int i=(b)-1;i>=(a);--i)
#define FILL(A,value) memset(A,value,sizeof(A))

#define ALL(V) V.begin(), V.end()
#define SZ(V) (int)V.size()
#define PB push_back
#define MP make_pair
#define Pi 3.14159265358979
#define x0 ikjnrmthklmnt
#define y0 lkrjhkltr
#define y1 ewrgrg

typedef long long Int;
typedef unsigned long long UInt;
typedef vector<int> VI;
typedef pair<int, int> PII;
typedef pair<Int, Int> PLL;
typedef pair<double, double> PDD;
typedef complex<double> base;

const int INF = 1000000000;
const int BASE = 1000000007;
const int MAX = 100007;
const int ADD = 1000000;
const int MOD = 1000000007;
const int CNT = 800;

int a[MAX];
VI G[MAX];
vector<Int> S[MAX];
vector<PII> Q[MAX];
Int R[MAX];

void dfs(int v, int dist , int p)
{
  S[v].push_back(a[v]);

  FOR(i,0,SZ(G[v]))
  {
    int to = G[v][i];
    if (to == p) continue;

    dfs(to,dist + 1, v);

    S[to].push_back(0);
    if (SZ(S[to]) > SZ(S[v])) S[v].swap(S[to]);

    FOR(j,0,SZ(S[to]))
    {
      S[v][j + SZ(S[v]) - SZ(S[to])] += S[to][j];
    }
    vector<Int> tmp;
    swap(tmp, S[to]);
  }

  FOR(i,0,SZ(Q[v]))
  {
    int d = Q[v][i].first - dist;
    int id = Q[v][i].second;

    if (d < 0) continue;
    if (d >= SZ(S[v])) continue;

    R[id] = S[v][SZ(S[v]) - 1 - d];
  }

}

int main()
{
    //freopen("in.txt", "r", stdin);
    //freopen("distance.in",  "r", stdin);
    //freopen("distance.out", "w", stdout);
    //freopen("out.txt" , "w" , stdout);

  int t;
  cin >> t;
  FOR(tt,0,t)
  {

    int n , q;
    cin >> n >> q;

    FOR(i,0,n)
    {
      G[i].clear();
      S[i].clear();
      Q[i].clear();
    }
    FOR(i,0,q)
    {
      R[i] = 0;
    }

    FOR(i,0,n)
    {
      scanf("%d" , &a[i]);
    }

    FOR(i,0,n - 1)
    {
      int u , v;
      scanf("%d%d" , &u , &v);
      -- u ; -- v;
      G[u].push_back(v);
      G[v].push_back(u);
    }

    FOR(i,0,q)
    {
      int u , l;
      scanf("%d%d" , &u , &l);
      -- u;
      -- l;
      Q[u].push_back(MP(l, i));
    }

    dfs(0,0,-1);

    FOR(i,0,q)
    {
      printf("%lld\n" , R[i]);
    }

  }


    return 0;
}