HackerEarth Christmas tree problem solution

In this HackerEarth Christmas tree problem solution Little Santa have to give gifts to N children. There are N - 1 roads connecting the homes of N children. Each road has some special Santa tax value which only Santa have to pay. Little Santa is wondering what is the realKth minimum tax value in the path between the home of child realA and child realB.

Note: kth minimum value in a list A is kth value in the list A after sorting A.

HackerEarth Christmas tree problem solution.

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>
#include <algorithm>
#include <set>
#include <map>
#include <cmath>
#include <ctime>
#include <functional>
#include <sstream>
#include <fstream>
#include <valarray>
#include <complex>
#include <queue>
#include <cassert>
#include <bitset>
using namespace std;

#ifdef LOCAL
#define debug_flag 1
#else
#define debug_flag 0
#endif

  template <class T1, class T2 >
std::ostream& operator << (std::ostream& os, const pair<T1, T2> &p) 
{
  os << "[" << p.first << ", " << p.second << "]";
  return os;
}

  template <class T >
std::ostream& operator << (std::ostream& os, const std::vector<T>& v) 
{
  os << "[";
  bool first = true;
  for (typename std::vector<T>::const_iterator it = v.begin(); it != v.end(); ++it)
  {
    if (!first)
      os << ", ";
    first = false;
    os << *it;
  }
  os << "]";
  return os;
}

#define dbg(args...) { if (debug_flag) { _print(_split(#args, ',').begin(), args); cerr << endl; } else { void(0);} }

vector<string> _split(const string& s, char c) {
  vector<string> v;
  stringstream ss(s);
  string x;
  while (getline(ss, x, c))
    v.emplace_back(x);
  return v;
}

void _print(vector<string>::iterator) {}
template<typename T, typename... Args>
void _print(vector<string>::iterator it, T a, Args... args) {
  string name = it -> substr((*it)[0] == ' ', it -> length());
  if (isalpha(name[0]))
    cerr << name  << " = " << a << " ";
  else
    cerr << name << " ";
  _print(++it, args...);
}

typedef long long int int64;

const int N = (int)4e5 + 100;
const int LOGN = 20;
const int pow2 = (1 << LOGN);

struct Edge
{
  int a, b, c;
  Edge() : a(), b(), c() {}
  Edge(int _a, int _b, int _c) : a(_a), b(_b), c(_c) {}

  bool operator < (const Edge & other) const
  {
    return c < other.c;
  }
};

struct Node
{
  int sum;
  Node *l, *r;

  Node()
  {

  }

  Node(int _sum)
  {
    sum = _sum;
    l = NULL;
    r = NULL;
  }

  Node(Node *_l, Node *_r)
  {
    l = _l;
    r = _r;
    sum = l->sum + r->sum;
  }
};

typedef Node * pNode;

const int B = 2 * N * LOGN;
int node_ptr;
Node nodes[B];

pNode newNode(int sum)
{
  assert(node_ptr < B);
  nodes[node_ptr].sum = sum;
  nodes[node_ptr].l = NULL;
  nodes[node_ptr].r = NULL;
  node_ptr++;
  return nodes + node_ptr - 1;
}

pNode newNode(Node *l, Node *r)
{
  assert(node_ptr < B);
  nodes[node_ptr].sum = l->sum + r->sum;
  nodes[node_ptr].l = l;
  nodes[node_ptr].r = r;
  node_ptr++;
  return nodes + node_ptr - 1;
}

pNode set_val(pNode v, int l, int r, int pos, int val)
{
  if (l == r)
    return newNode(val);
  int m = (l + r) / 2;
  if (pos <= m)
    return newNode(set_val(v->l, l, m, pos, val), v->r);
  return newNode(v->l, set_val(v->r, m + 1, r, pos, val));
}

int n;
vector<pair<int, int> > graph[N];
Edge e_list[N];

pNode tree[N];

int timer;
int t_in[N], t_out[N];

int par[N][LOGN];

void scan_graph()
{
  scanf("%d", &n);

  for (int i = 0; i < n - 1; i++)
  {
    int a, b, c;
    scanf("%d%d%d", &a, &b, &c);
    a--;
    b--;
    e_list[i] = Edge(a, b, c);
  }

  sort(e_list, e_list + n - 1);

  for (int i = 0; i < n - 1; i++)
  {
    Edge e = e_list[i];
    graph[e.a].emplace_back(e.b, i);
    graph[e.b].emplace_back(e.a, i);
  }
}

int get_ans(pNode t1, pNode t2, pNode t3, int l, int r, int k)
{
  if (l == r)
    return l;

  int m = (l + r) / 2;
  int left_sum = t1->l->sum + t2->l->sum - 2 * t3->l->sum;

  if (k < left_sum)
    return get_ans(t1->l, t2->l, t3->l, l, m, k);

  return get_ans(t1->r, t2->r, t3->r, m + 1, r, k - left_sum);
}

int get_ans(pNode t1, pNode t2, pNode t3, int k)
{
  if (k >= t1->sum + t2->sum - 2 * t3->sum)
    return -1;
  return e_list[get_ans(t1, t2, t3, 0, pow2 - 1, k)].c;
}

void precalc(int v, int p, pNode t)
{
  t_in[v] = timer++;
  tree[v] = t;

  par[v][0] = p;
  for (int i = 1; i < LOGN; i++)
  {
    if (par[v][i - 1] == -1)
      break;
    par[v][i] = par[par[v][i - 1]][i - 1];
  }

  for (auto pp : graph[v])
  {
    int to = pp.first;
    int w = pp.second;
    if (to == p)
      continue;
    precalc(to, v, set_val(t, 0, pow2 - 1, w, 1));
  }

  t_out[v] = timer++;
}

bool is_par(int p, int v)
{
  return t_in[p] <= t_in[v] && t_out[v] <= t_out[p];
}

int get_lca(int a, int b)
{
  if (is_par(a, b))
    return a;
  if (is_par(b, a))
    return b;
  for (int i = LOGN - 1; i >= 0; i--)
  {
    int new_a = par[a][i];
    if (new_a == -1 || is_par(new_a, b))
      continue;
    a = new_a;
  }
  return par[a][0];
}

int solve(int a, int b, int k)
{
  int lca = get_lca(a, b);
  return get_ans(tree[a], tree[b], tree[lca], k);
}

int main()
{
#ifdef LOCAL
  freopen ("input.txt", "r", stdin);
#endif

  scan_graph();
  
  pNode empty_tree = newNode(0);
  for (int i = 0; i < LOGN; i++)
    empty_tree = newNode(empty_tree, empty_tree);

  precalc(0, -1, empty_tree);

  int q;
  scanf("%d", &q);
  int64 last_ans = 0;

  for (int i = 0; i < q; i++)
  {
    int a, b, k;
    scanf("%d%d%d", &a, &b, &k);
  
    a = (a - 1 + max(0LL, last_ans)) % n;
    b = (b - 1 + max(0LL, 2 * last_ans)) % n;
    k = (k - 1 + max(0LL, 3 * last_ans)) % n;

    int ans = solve(a, b, k);
    printf("%d\n", ans);
    last_ans = ans;
  }

  return 0;
}