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("%lldn", 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("%lldn" , R[i]);
}
}
return 0;
}