Outline

Introduction

Implementation

int n, q, cnt, sz[N], dep[N], fa[N], son[N], pos[N], head[N];
vector<int> G[N];

void dfs(int u, int d, int f) {
  sz[u] = 1, dep[u] = d, fa[u] = f;
  int mx = 0;
  for(auto v : G[u]) if(v != f) {
    dfs(v, d + 1, u);
    sz[u] += sz[v];
    if(sz[v] > mx) mx = sz[v], son[u] = v;
  }
}
void hld(int u, int F) {
  pos[u] = cnt++, head[u] = F;
  for(auto v : G[u]) if(v == son[u] && v != fa[u]) {
    hld(v, F);
    break;
  }
  for(auto v : G[u]) if(v != son[u] && v != fa[u]) {
    hld(v, v);
  }
}

Example problems

SPOJ - QTREE

Problem description

You are given a weighted tree consist of $N$ nodes, with edges indexed from $1$ to $N-1$, and you need to handle $Q$ queries of two kinds:

  • $CHANGE\space i\space t_i$: change the $i^{th}$ edges weight to $t_i$.
  • $QUERY\space a\space b$: Answer the maximum edge weight on the path from $a$ to $b$.

$2\le N \le 10^4$

Problem analysis

Problem solution

code 
#include <bits/stdc++.h>
#define F first
#define S second
using namespace std;
using ll = long long;

const int N = (int)1e4 + 5;
const int inf = (int)1e9;

namespace segtree {
  int mx[N * 4], a[N];
  inline void pull(int o) {
    mx[o] = max(mx[o * 2 + 1], mx[o * 2 + 2]);
  }
  void build(int l, int r, int o=0) {
    if(l + 1 == r) return void(mx[o] = a[l]);
    int mid = (l + r) >> 1;
    build(l, mid, o * 2 + 1); build(mid, r, o * 2 + 2);
    pull(o);
  }
  void update(int l, int r, int q, int v, int o=0) {
    if(l + 1 == r) return void(mx[o] = v);
    int mid = (l + r) >> 1;
    if(q < mid) update(l, mid, q, v, o * 2 + 1);
    else update(mid, r, q, v, o * 2 + 2);
    pull(o);
  }
  int query(int l, int r, int ql, int qr, int o=0) {
    if(ql <= l && r <= qr) return mx[o];
    int mid = (l + r) >> 1;
    if(qr <= mid) return query(l, mid, ql, qr, o * 2 + 1);
    else if(mid <= ql) return query(mid, r, ql, qr, o * 2 + 2);
    return max(query(l, mid, ql, mid, o * 2 + 1), query(mid, r, mid, qr, o * 2 + 2));
  }
}

int n, cnt, dep[N], pos[N], sz[N], son[N], fa[N], head[N]; // son = heavy son
vector<pair<int, int> > G[N], edges;

void dfs(int u, int d, int f) {
  sz[u] = 1, fa[u] = f, dep[u] = d;
  int mx = 0;
  for(auto &e : G[u]) if(e.F != f) {
    dfs(e.F, d + 1, u);
    sz[u] += sz[e.F];
    if(sz[e.F] > mx) mx = sz[e.F], son[u] = e.F;
  }
}
void hld(int u, int w, int F) {
  head[u] = F, pos[u] = cnt;
  segtree::a[cnt++] = w;
  for(auto &e : G[u]) if(e.F == son[u] && e.F != son[u]) {
    hld(e.F, e.S, F);
    break;
  }
  for(auto &e : G[u]) if(e.F != fa[u] && e.F != son[u]) {
    hld(e.F, e.S, e.F);
  }
}
int query(int u, int v) {
  int mx = -inf;
  while(head[u] != head[v]) {
    if(dep[head[u]] > dep[head[v]]) swap(u, v);
    mx = max(mx, segtree::query(0, n, pos[head[v]], pos[v] + 1));
    v = fa[head[v]];
  }
  if(dep[u] > dep[v]) swap(u, v);
  if(pos[u] < pos[v]) mx = max(mx, segtree::query(0, n, pos[u] + 1, pos[v] + 1));
  return mx;
}
void update(int i, int w) {
  int u = dep[edges[i].F] > dep[edges[i].S] ? edges[i].F : edges[i].S;
  segtree::update(0, n, pos[u], w);
}

void init() {
  cin >> n;
  cnt = 0;
  for(int i = 0 ; i < n ; ++i) G[i].clear();
  edges.clear();
  for(int i = 0 ; i < n - 1 ; ++i) {
    int a, b, c; cin >> a >> b >> c; a--, b--;
    G[a].push_back({b, c});
    G[b].push_back({a, c});
    edges.push_back({a, b});
  }
}
void solve() {
  dfs(0, 0, 0);
  hld(0, -inf, 0);
  segtree::build(0, n);
  string s;
  while(cin >> s && s[0] != 'D') {
    int i, u, v, w;
    if(s[0] == 'Q') cin >> u >> v, cout << query(u - 1, v - 1) << '\n';
    else cin >> i >> w, update(i - 1, w);
  }
} 

int main() {
  ios_base::sync_with_stdio(0), cin.tie(0);
  int t; cin >> t;
  while(t--) {
    init();
    solve();
  }
}

SPOJ - QTREE2

Problem description

You are given a weighted tree consist of $N$ nodes, with edges indexed from $1$ to $N-1$, and you need to handle $Q$ queries of two kinds:

  • $DIST\space a\space b$: Answer the distance between $a$ and $b$.
  • $KTH\space a\space b\space k$: Answer the $k^{th}$ node on the path from $a$ to $b$..

$2\le N \le 10^4$

Problem analysis

Problem solution

code 
#include <bits/stdc++.h>
#define F first
#define S second
using namespace std;
using ll = long long;

const int N = (int)1e4 + 5;

int n, cnt, cnt2, sz[N], dep[N], fa[N], son[N], head[N];
pair<int, int> pos[N];
vector<pair<int, int> > G[N];
vector<int> node[N], dis[N];

void dfs(int u, int d, int f) {
  sz[u] = 1, dep[u] = d, fa[u] = f;
  int mx = 0;
  for(auto &e : G[u]) if(e.F != f) {
    dfs(e.F, d + 1, u);
    sz[u] += sz[e.F];
    if(sz[e.F] > mx) mx = sz[e.F], son[u] = e.F;
  }
}
void hld(int u, int w, int F) {
  node[cnt].push_back(u);
  dis[cnt].push_back(w);
  pos[u] = {cnt, cnt2++};
  head[u] = F;
  for(auto &e : G[u]) if(e.F == son[u] && e.F != fa[u]) {
    hld(e.F, e.S, F);
    break;
  }
  for(auto &e : G[u]) if(e.F != fa[u] && e.F != son[u]) {
    cnt++; cnt2 = 0;
    hld(e.F, e.S, e.F);
  }
}

int get_dis(pair<int, int> p1, pair<int, int> p2) {
  assert(p1.F == p2.F);
  if(p1.S == 0) return dis[p2.F][p2.S];
  return dis[p2.F][p2.S] - dis[p1.F][p1.S - 1];
}
int query1(int u, int v) {
  int sum = 0;
  while(head[u] != head[v]) {
    if(dep[head[u]] > dep[head[v]]) swap(u, v);
    sum += get_dis(pos[head[v]], pos[v]);
    v = fa[head[v]];
  }
  if(dep[u] > dep[v]) swap(u, v);
  if(pos[u] < pos[v]) sum += get_dis({pos[u].F, pos[u].S + 1}, pos[v]);
  return sum;
}
int query2(int u, int v, int k) {
  int len = 0, _u = u, _v = v;
  while(head[u] != head[v]) {
    if(dep[head[u]] > dep[head[v]]) swap(u, v);
    len += pos[v].S - pos[head[v]].S + 1;
    v = fa[head[v]];
  }
  if(dep[u] > dep[v]) swap(u, v);
  len += pos[v].S - pos[u].S + 1;
  assert(k <= len);
  u = _u, v = _v;
  while(head[u] != head[v]) {
    if(dep[head[u]] > dep[head[v]]) {
      int _k = k;
      k -= pos[u].S - pos[head[u]].S + 1;
      len -= _k - k;
      if(k < 1) {
        return node[pos[u].F][pos[u].S - _k + 1];
      }
      u = fa[head[u]];
    }
    else {
      int _len = len;
      len -= pos[v].S - pos[head[v]].S + 1;
      if(len < k) {
        return node[pos[v].F][pos[v].S - (_len - k)];
      }
      v = fa[head[v]];
    }
  }
  if(dep[u] > dep[v]) return node[pos[u].F][pos[u].S - k + 1];
  else return node[pos[u].F][pos[u].S + k - 1];
}

void init() {
  for(int i = 0 ; i <= cnt ; ++i) node[i].clear(), dis[i].clear();
  cin >> n;
  for(int i = 0 ; i < n ; ++i) G[i].clear();
  for(int i = 0 ; i < n - 1 ; ++i) {
    int a, b, c; cin >> a >> b >> c; a--, b--;
    G[a].push_back({b, c});
    G[b].push_back({a, c});
  }
}
void solve() {
  dfs(0, 0, 0);
  cnt = cnt2 = 0;
  hld(0, 0, 0);
  for(int i = 0 ; i <= cnt ; ++i) {
    for(int j = 1 ; j < (int)dis[i].size() ; ++j) dis[i][j] += dis[i][j - 1];
  }
  string s;
  while(cin >> s && s[1] != 'O') {
    int a, b, k;
    if(s[1] == 'I') {
      cin >> a >> b; a--, b--;
      cout << query1(a, b) << '\n';
    }
    else {
      cin >> a >> b >> k; a--, b--;
      cout << query2(a, b, k) + 1 << '\n';
    }
  }
}

int main() {
  ios_base::sync_with_stdio(0), cin.tie(0);
  int t; cin >> t;
  while(t--) {
    init();
    solve();
  }
}

SPOJ - QTREE6

Problem description

You are given a tree consist of $N$ nodes, with nodes indexed from $1$ to $N$, and you need to handle $Q$ queries of two kinds:

  • $0\space u$: Answer the number of nodes connected to $u$. Two nodes are connected iff all nodes on path between them have same color.
  • $1\space u$: Change the color of node $u$.

$1\le N, Q \le 10^5$

Problem analysis

Problem solution

code 
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;

const int N = (int)1e5 + 5;
const int inf = (int)1e9;

struct Segtree {
  int b[N], a[N * 4];
  void build(int l, int r, int o=0) {
    if(l + 1 == r) return void(a[o] = b[l]);
    int mid = (l + r) >> 1;
    build(l, mid, o * 2 + 1);
    build(mid, r, o * 2 + 2);
  }
  void update(int l, int r, int ql, int qr, int v, int o=0) {
    if(ql <= l && r <= qr) return void(a[o] += v);
    int mid = (l + r) >> 1;
    if(qr <= mid) update(l, mid, ql, qr, v, o * 2 + 1);
    else if(mid <= ql) update(mid, r, ql, qr, v, o * 2 + 2);
    else {
      update(l, mid, ql, mid, v, o * 2 + 1);
      update(mid, r, mid, qr, v, o * 2 + 2);
    }
  }
  int query(int l, int r, int q, int o=0) {
    if(l + 1 == r) return a[o];
    int mid = (l + r) >> 1;
    if(q < mid) return a[o] + query(l, mid, q, o * 2 + 1);
    else return a[o] + query(mid, r, q, o * 2 + 2);
  }
} seg_sz[2];

struct Segtree2{
  int b[N], a[N * 4];
  void build(int l, int r, int o=0) {
    if(l + 1 == r) return void(a[o] = (b[l] ? l : -inf));
    int mid = (l + r) >> 1;
    build(l, mid, o * 2 + 1);
    build(mid, r, o * 2 + 2);
    a[o] = max(a[o * 2 + 1], a[o * 2 + 2]);
  }
  void update(int l, int r, int q, int v, int o=0) {
    if(l + 1 == r) return void(a[o] = (v ? l : -inf));
    int mid = (l + r) >> 1;
    if(q < mid) update(l, mid, q, v, o * 2 + 1);
    else update(mid, r, q, v, o * 2 + 2);
    a[o] = max(a[o * 2 + 1], a[o * 2 + 2]);
  }
  int query(int l, int r, int ql, int qr, int o=0) {
    if(ql <= l && r <= qr) return a[o];
    int mid = (l + r) >> 1;
    if(qr <= mid) return query(l, mid, ql, qr, o * 2 + 1);
    if(mid <= ql) return query(mid, r, ql, qr, o * 2 + 2);
    return max(query(l, mid, ql, mid, o * 2 + 1), query(mid, r, mid, qr, o * 2 + 2));
  }
} seg_pos[2];

int n, cnt = 0, sz[N], fa[N], dep[N], son[N], head[N], pos[N], color[N];
vector<int> G[N];

void dfs(int u, int d, int f) {
  sz[u] = 1, dep[u] = d, fa[u] = f;
  int mx = 0;
  for(auto v : G[u]) if(v != f) {
    dfs(v, d + 1, u);
    sz[u] += sz[v];
    if(sz[v] > mx) mx = sz[v], son[u] = v;
  }
}
void hld(int u, int F) {
  seg_sz[0].b[cnt] = sz[u], seg_sz[1].b[cnt] = 1;
  seg_pos[0].b[cnt] = 1, seg_pos[1].b[cnt] = 0;
  head[u] = F, pos[u] = cnt++;
  for(auto v : G[u]) if(v != fa[u] && v == son[u]) {
    hld(v, F);
    break;
  }
  for(auto v : G[u]) if(v != fa[u] && v != son[u]) {
    hld(v, v);
  }
}
void update(int u) {
  // add to new color
  int v = fa[u], size = seg_sz[color[u] ^ 1].query(0, n, pos[u]);
  while(v != -1) {
    int low_pos = seg_pos[color[u]].query(0, n, pos[head[v]], pos[v] + 1);
    seg_sz[color[u] ^ 1].update(0, n, low_pos == -inf ? pos[head[v]] : low_pos, pos[v] + 1, size);
    if(low_pos != -inf) break;
    v = fa[head[v]];
  }
  seg_pos[color[u]].update(0, n, pos[u], 0);
  // change color
  color[u] ^= 1;
  // subtract from old color
  seg_pos[color[u]].update(0, n, pos[u], 1);
  v = fa[u], size = seg_sz[color[u] ^ 1].query(0, n, pos[u]);
  while(v != -1) {
    int low_pos = seg_pos[color[u]].query(0, n, pos[head[v]], pos[v] + 1);
    seg_sz[color[u] ^ 1].update(0, n, low_pos == -inf ? pos[head[v]] : low_pos, pos[v] + 1, -size);
    if(low_pos != -inf) break;
    v = fa[head[v]];
  }
}
int query(int u) {
  int v = u, low_pos;
  while(v != -1) {
    low_pos = seg_pos[color[u] ^ 1].query(0, n, pos[head[v]], pos[v] + 1);
    if(low_pos != -inf) break;
    u = head[v], v = fa[head[v]];
  }
  if(low_pos == -inf) low_pos = 0;
  else {
    if(low_pos == pos[v]) low_pos = pos[u];
    else low_pos += 1;
  }
  return seg_sz[color[u]].query(0, n, low_pos); 
}

void init() {
  cin >> n;
  for(int i = 0 ; i < n - 1 ; ++i) {
    int u, v; cin >> u >> v; --u, --v;
    G[u].push_back(v);
    G[v].push_back(u);
  }
  dfs(0, 0, -1);
  hld(0, 0);
  seg_sz[0].build(0, n); seg_sz[1].build(0, n);
  seg_pos[0].build(0, n); seg_pos[1].build(0, n);
}
void solve() {
  int q; cin >> q;
  while(q--) {
    int op, u; cin >> op >> u; --u;
    if(op == 0) cout << query(u) << '\n';
    else update(u);
  }
}

int main() {
  ios_base::sync_with_stdio(0), cin.tie(0);
  init();
  solve();
}

Codechef - DGCD

Problem description

You are given a tree consist of $N$ nodes. Each node has a positive integer written on it, number on the $i^{th}$ vertex being $v_i$. You need to handle $Q$ queries of two kinds:

  • $F\space u\space v$: Find out the $gcd$(link) of all numbers on the unique path between $u$ and $v$.
  • $C\space u\space v\space d$: Add $d$ to the number written on all nodes along the path between $u$ and $v$ (both inclusive).

$1\le N, Q\le 5\times10^4, 0\le u, v, \le N - 1, 1\le v_i\le 10^4, 0\le d\le 10^4$

Problem analysis

Problem solution

code 
#include <bits/stdc++.h>
using namespace std;
using ll = long long;

const int N = (int)5e4 + 5;

namespace segtree {
  int a[N], raw[N * 4], dif[N * 4];

  void build_raw(int l, int r, int o=0) {
    if(l + 1 == r) return void(raw[o] = a[l]);
    int mid = (l + r) >> 1;
    build_raw(l, mid, o * 2 + 1); build_raw(mid, r, o * 2 + 2);
    raw[o] = 0;
  }
  inline void pull_dif(int o) {
    dif[o] = __gcd(dif[o * 2 + 1], dif[o * 2 + 2]);
  }
  void build_dif(int l, int r, int o=0) {
    if(l + 1 == r) return void(dif[o] = a[l] - a[l - 1]);
    int mid = (l + r) >> 1;
    build_dif(l, mid, o * 2 + 1); build_dif(mid, r, o * 2 + 2);
    pull_dif(o);
  }
  void build(int l, int r) {
    build_raw(l, r);
    build_dif(l + 1, r);
  }
  void update_raw(int l, int r, int ql, int qr, int v, int o=0) {
    if(ql <= l && r <= qr) return void(raw[o] += v);
    int mid = (l + r) >> 1;
    if(qr <= mid) update_raw(l, mid, ql, qr, v, o * 2 + 1);
    else if(mid <= ql) update_raw(mid, r, ql, qr, v, o * 2 + 2);
    else {
      update_raw(l, mid, ql, mid, v, o * 2 + 1);
      update_raw(mid, r, mid, qr, v, o * 2 + 2);
    }
  }
  void update_dif(int l, int r, int q, int v, int o=0) {
    if(l + 1 == r) return void(dif[o] += v);
    int mid = (l + r) >> 1;
    if(q < mid) update_dif(l, mid, q, v, o * 2 + 1);
    else update_dif(mid, r, q, v, o * 2 + 2);
    pull_dif(o);
  }
  void update(int l, int r, int ql, int qr, int v) {
    update_raw(l, r, ql, qr, v);
    if(ql > 0) update_dif(l + 1, r, ql, v);
    if(qr < r) update_dif(l + 1, r, qr, -v);
  }
  int query_raw(int l, int r, int q, int o=0) {
    if(l + 1 == r) return raw[o];
    int mid = (l + r) >> 1;
    if(q < mid) return raw[o] + query_raw(l, mid, q, o * 2 + 1);
    return raw[o] + query_raw(mid, r, q, o * 2 + 2);
  }
  int query_dif(int l, int r, int ql, int qr, int o=0) {
    if(ql <= l && r <= qr) return dif[o];
    int mid = (l + r) >> 1;
    if(qr <= mid) return query_dif(l, mid, ql, qr, o * 2 + 1);
    else if(mid <= ql) return query_dif(mid, r, ql, qr, o * 2 + 2);
    else return __gcd(query_dif(l, mid, ql, mid, o * 2 + 1), query_dif(mid, r, mid, qr, o * 2 + 2));
  }
  int query(int l, int r, int ql, int qr) {
    if(ql + 1 >= qr) return query_raw(l, r, ql);
    return __gcd(query_raw(l, r, ql), query_dif(l + 1, r, ql + 1, qr));
  }
}

int n, cnt = 0, num[N], sz[N], fa[N], son[N], head[N], pos[N], dep[N];
vector<int> G[N];

void dfs(int u, int d, int f) {
  sz[u] = 1, dep[u] = d, fa[u] = f;
  int mx = 0;
  for(auto v : G[u]) if(v != f) {
    dfs(v, d + 1, u);
    sz[u] += sz[v];
    if(sz[v] > mx) mx = sz[v], son[u] = v;
  }
}
void hld(int u, int F) {
  head[u] = F, pos[u] = cnt;
  segtree::a[cnt++] = num[u];
  for(auto v : G[u]) if(v != fa[u] && v == son[u]) {
    hld(v, F);
    break;
  }
  for(auto v : G[u]) if(v != fa[u] && v != son[u]) {
    hld(v, v);
  }
}
int query(int u, int v) {
  int ans = 0;
  while(head[u] != head[v]) {
    if(dep[head[u]] > dep[head[v]]) swap(u, v);
    int val = segtree::query(0, n, pos[head[v]], pos[v] + 1);
    ans = (ans == 0 ? val : __gcd(ans, val));
    v = fa[head[v]];
  }
  if(dep[u] > dep[v]) swap(u, v);
  int val = segtree::query(0, n, pos[u], pos[v] + 1);
  ans = (ans == 0 ? val : __gcd(ans, val));
  return abs(ans);
}
void update(int u, int v, int d) {
  while(head[u] != head[v]) {
    if(dep[head[u]] > dep[head[v]]) swap(u, v);
    segtree::update(0, n, pos[head[v]], pos[v] + 1, d);
    v = fa[head[v]];
  }
  if(dep[u] > dep[v]) swap(u, v);
  segtree::update(0, n, pos[u], pos[v] + 1, d);
}

void init() {
  cin >> n;
  for(int i = 0 ; i < n - 1 ; ++i) {
    int a, b; cin >> a >> b;
    G[a].push_back(b);
    G[b].push_back(a);
  }
  for(int i = 0 ; i < n ; ++i) cin >> num[i];
}
void solve() {
  dfs(0, 0, 0);
  hld(0, 0);
  segtree::build(0, n);
  int q; cin >> q;
  while(q--) {
    string op; cin >> op;
    if(op[0] == 'F') {
      int u, v; cin >> u >> v;
      cout << query(u, v) << '\n';
    }
    else {
      int u, v, d; cin >> u >> v >> d;
      update(u, v, d);
    }
  }
}

int main() {
  ios_base::sync_with_stdio(0), cin.tie(0);
  init();
  solve();
}

Codechef - QUERY

Problem description

Problem analysis

Problem solution

More problems