SubtreesAndPaths.java

package io.github.chermehdi.main;

import io.github.chermehdi.lib.ArrayUtils;
import io.github.chermehdi.lib.Bits;
import io.github.chermehdi.lib.GraphUtils;
import io.github.chermehdi.lib.externals.SegmentTree;
import io.github.chermehdi.lib.io.InputReader;
import io.github.chermehdi.lib.io.OutputWriter;
import java.util.List;

public class SubtreesAndPaths {

  int[] depth, parent, head, heavy, next, size;
  int[] first, last;
  int n;
  List<Integer>[] g;
  int[][] sparse;
  int index = 0;
  SegmentTree segTree;

  void init() {
    depth = new int[n + 1];
    parent = new int[n + 1];
    head = new int[n + 1];
    heavy = new int[n + 1];
    size = new int[n + 1];
    next = new int[n + 1];
    first = ArrayUtils.createArray(n + 1, -1);
    last = new int[n + 1];
    sparse = new int[20][n + 1];
    segTree = new SegmentTree(n + 1);
  }

  void dfs(int u, int p) {
    size[u] = 1;
    parent[u] = p;
    sparse[0][u] = p;
    int maxSize = 0;
    for (int v : g[u]) {
      if (v != p) {
        depth[v] = depth[u] + 1;
        dfs(v, u);
        size[u] += size[v];
        if (size[v] > maxSize) {
          maxSize = size[v];
          heavy[u] = v;
        }
      }
    }
  }

  int query(int u, int v) {
    int res = -(1 << 30);
    while (next[u] != next[v]) {
      if (depth[next[u]] > depth[next[v]]) {
        int temp = u;
        u = v;
        v = temp;
      }
      int cur = segTree.query(first[next[v]], first[v]);
      res = Math.max(res, cur);
      v = parent[next[v]];
    }
    if (depth[u] > depth[v]) {
      int temp = u;
      u = v;
      v = temp;
    }
    return Math.max(res, segTree.query(first[u], first[v]));
  }

  void decompose(int u, int p) {
    first[u] = index++;
    for (int v : g[u]) {
      if (v != p) {
        next[v] = (v == heavy[u] ? next[u] : v);
        decompose(v, u);
      }
    }
    last[u] = index;
  }


  int lca(int u, int v) {
    if (depth[u] < depth[v]) {
      int temp = u;
      u = v;
      v = temp;
    }
    int diff = depth[u] - depth[v];
    for (int i = 0; i < 20; ++i) {
      if (Bits.isSet(diff, i)) {
        u = sparse[i][u];
      }
    }
    if (u == v) {
      return u;
    }
    for (int i = 19; i >= 0; --i) {
      if (sparse[i][u] != sparse[i][v]) {
        u = sparse[i][u];
        v = sparse[i][v];
      }
    }
    return sparse[0][u];
  }

  public void solve(int testNumber, InputReader in, OutputWriter out) {
    n = in.nextInt();
    g = GraphUtils.generateGraph(n + 1);
    for (int i = 0; i < n - 1; ++i) {
      int u = in.nextInt();
      int v = in.nextInt();
      g[u].add(v);
      g[v].add(u);
    }
    init();
    // suppose root is 1
    dfs(1, 0);

    for (int i = 1; i < 20; ++i) {
      for (int u = 1; u <= n; ++u) {
        sparse[i][u] = sparse[i - 1][sparse[i - 1][u]];
      }
    }

    int q = in.nextInt();
    for (int i = 1; i <= n; ++i) {
      next[i] = i;
    }
    decompose(1, 1);

    while (q-- > 0) {
      String type = in.next();
      if (type.charAt(0) == 'a') {
        int node = in.nextInt();
        int val = in.nextInt();
        segTree.modify(first[node], last[node] - 1, val);
      } else {
        int u = in.nextInt();
        int v = in.nextInt();
        int p = lca(u, v);
        out.println(Math.max(query(p, u), query(p, v)));
      }
    }
  }

  void modify(int u, int p, int val) {
    segTree.modify(first[u], first[u], val);
    for (int v : g[u]) {
      if (v != p) {
        modify(v, u, val);
      }
    }
  }
}