Heavy-light tree decomposition for vertices or edges

import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;

// Heavy-light decomposition with path queries. Query complexity is O(log^2(n)).
// Based on the code from http://codeforces.com/blog/entry/22072
public class HeavyLight {

  int getNeutralValue() {
    return 0;
  }

  final List<Integer>[] tree;
  final boolean valuesOnVertices; // true - values on vertices, false - values on edges
  final SegmentTree segmentTree;
  final int[] parent;
  final int[] heavy;
  final int[] depth;
  final int[] pathRoot;
  final int[] pos;

  public HeavyLight(List<Integer>[] tree, boolean valuesOnVertices) {
    this.tree = tree;
    this.valuesOnVertices = valuesOnVertices;
    int n = tree.length;
    segmentTree = new SegmentTree(n);

    parent = new int[n];
    heavy = new int[n];
    depth = new int[n];
    pathRoot = new int[n];
    pos = new int[n];

    Arrays.fill(heavy, -1);
    parent[0= -1;
    depth[00;
    dfs(0);
    for (int u = 0, p = 0; u < n; u++) {
      if (parent[u== -|| heavy[parent[u]] != u) {
        for (int v = u; v != -1; v = heavy[v]) {
          pathRoot[v= u;
          pos[v= p++;
        }
      }
    }
  }

  int dfs(int u) {
    int size = 1;
    int maxSubtree = 0;
    for (int v : tree[u]) {
      if (v != parent[u]) {
        parent[v= u;
        depth[v= depth[u1;
        int subtree = dfs(v);
        if (maxSubtree < subtree) {
          maxSubtree = subtree;
          heavy[u= v;
        }
        size += subtree;
      }
    }
    return size;
  }

  public int query(int u, int v) {
    AtomicInteger res = new AtomicInteger(getNeutralValue())// just mutable integer
    processPath(u, v, (a, b-> res.set(segmentTree.queryOperation(res.get(), segmentTree.query(a, b))));
    return res.get();
  }

  public void modify(int u, int v, int delta) {
    processPath(u, v, (a, b-> segmentTree.modify(a, b, delta));
  }

  void processPath(int u, int v, BiConsumer<Integer, Integer> op) {
    for (; pathRoot[u!= pathRoot[v]; v = parent[pathRoot[v]]) {
      if (depth[pathRoot[u]] > depth[pathRoot[v]]) {
        int t = u;
        u = v;
        v = t;
      }
      op.accept(pos[pathRoot[v]], pos[v]);
    }
    if (!valuesOnVertices && u == vreturn;
    op.accept(Math.min(pos[u], pos[v]) (valuesOnVertices ? 1), Math.max(pos[u], pos[v]));
  }

  static class SegmentTree {
    // Modify the following 5 methods to implement your custom operations on the tree.
    // This example implements Add/Sum operations. Operations like Add/Max, Set/Max can also be implemented.
    int modifyOperation(int x, int y) {
      return x + y;
    }

    // query (or combine) operation
    int queryOperation(int leftValue, int rightValue) {
      return leftValue + rightValue;
    }

    int deltaEffectOnSegment(int delta, int segmentLength) {
      if (delta == getNeutralDelta()) return getNeutralDelta();
      // Here you must write a fast equivalent of following slow code:
      // int result = delta;
      // for (int i = 1; i < segmentLength; i++) result = queryOperation(result, delta);
      // return result;
      return delta * segmentLength;
    }

    int getNeutralDelta() {
      return 0;
    }

    int getInitValue() {
      return 0;
    }

    // generic code
    int[] value;
    int[] delta; // delta[i] affects value[i], delta[2*i+1] and delta[2*i+2]

    int joinValueWithDelta(int value, int delta) {
      if (delta == getNeutralDelta()) return value;
      return modifyOperation(value, delta);
    }

    int joinDeltas(int delta1, int delta2) {
      if (delta1 == getNeutralDelta()) return delta2;
      if (delta2 == getNeutralDelta()) return delta1;
      return modifyOperation(delta1, delta2);
    }

    void pushDelta(int i) {
      int d = 0;
      for ((i >> d0; d++) {
      }
      for (d -= 2; d >= 0; d--) {
        int x = i >> d;
        value[x >> 1= joinNodeValueWithDelta(x >> 1<< (d + 1));
        delta[x= joinDeltas(delta[x], delta[x >> 1]);
        delta[x ^ 1= joinDeltas(delta[x ^ 1], delta[x >> 1]);
        delta[x >> 1= getNeutralDelta();
      }
    }

    public SegmentTree(int n) {
      value = new int[* n];
      for (int i = 0; i < n; i++) {
        value[i + n= getInitValue();
      }
      for (int i = * n - 1; i > 1; i -= 2) {
        value[i >> 1= queryOperation(value[i], value[i ^ 1]);
      }
      delta = new int[* n];
      Arrays.fill(delta, getNeutralDelta());
    }

    int joinNodeValueWithDelta(int i, int len) {
      return joinValueWithDelta(value[i], deltaEffectOnSegment(delta[i], len));
    }

    public int query(int from, int to) {
      from += value.length >> 1;
      to += value.length >> 1;
      pushDelta(from);
      pushDelta(to);
      int res = 0;
      boolean found = false;
      for (int len = 1; from <= to; from = (from + 1>> 1, to = (to - 1>> 1, len <<= 1) {
        if ((from & 1!= 0) {
          res = found ? queryOperation(res, joinNodeValueWithDelta(from, len)) : joinNodeValueWithDelta(from, len);
          found = true;
        }
        if ((to & 1== 0) {
          res = found ? queryOperation(res, joinNodeValueWithDelta(to, len)) : joinNodeValueWithDelta(to, len);
          found = true;
        }
      }
      if (!foundthrow new RuntimeException();
      return res;
    }

    public void modify(int from, int to, int delta) {
      from += value.length >> 1;
      to += value.length >> 1;
      pushDelta(from);
      pushDelta(to);
      int a = from;
      int b = to;
      for (; from <= to; from = (from + 1>> 1, to = (to - 1>> 1) {
        if ((from & 1!= 0) {
          this.delta[from= joinDeltas(this.delta[from], delta);
        }
        if ((to & 1== 0) {
          this.delta[to= joinDeltas(this.delta[to], delta);
        }
      }
      for (int i = a, len = 1; i > 1; i >>= 1, len <<= 1) {
        value[i >> 1= queryOperation(joinNodeValueWithDelta(i, len), joinNodeValueWithDelta(i ^ 1, len));
      }
      for (int i = b, len = 1; i > 1; i >>= 1, len <<= 1) {
        value[i >> 1= queryOperation(joinNodeValueWithDelta(i, len), joinNodeValueWithDelta(i ^ 1, len));
      }
    }
  }

  // Random test
  public static void main(String[] args) {
    Random rnd = new Random(1);
    for (int step = 0; step < 1000; step++) {
      int n = rnd.nextInt(501;
      List<Integer>[] tree = getRandomTree(n, rnd);
      HeavyLight hl = new HeavyLight(tree, true);
      int[] x = new int[n];
      Arrays.fill(x, hl.segmentTree.getInitValue());
      for (int i = 0; i < 1000; i++) {
        int a = rnd.nextInt(n);
        int b = rnd.nextInt(n);
        List<Integer> path = new ArrayList<>();
        getPathFromAtoB(tree, a, b, -1, path);
        if (rnd.nextBoolean()) {
          int delta = rnd.nextInt(50100;
          hl.modify(a, b, delta);
          for (int u : path)
            x[u= hl.segmentTree.joinValueWithDelta(x[u], delta);
        else {
          int res1 = hl.query(a, b);
          int res2 = hl.getNeutralValue();
          for (int u : path)
            res2 = hl.segmentTree.queryOperation(res2, x[u]);
          if (res1 != res2)
            throw new RuntimeException();
        }
      }
    }

    for (int step = 0; step < 1000; step++) {
      int n = rnd.nextInt(501;
      List<Integer>[] tree = getRandomTree(n, rnd);
      HeavyLight hl = new HeavyLight(tree, false);
      Map<Long, Integer> x = new HashMap<>();
      for (int u = 0; u < tree.length; u++)
        for (int v : tree[u])
          x.put(edge(u, v), hl.segmentTree.getInitValue());
      for (int i = 0; i < 1000; i++) {
        int a = rnd.nextInt(n);
        int b = rnd.nextInt(n);
        List<Integer> path = new ArrayList<>();
        getPathFromAtoB(tree, a, b, -1, path);
        if (rnd.nextBoolean()) {
          int delta = rnd.nextInt(50100;
          hl.modify(a, b, delta);
          for (int j = 0; j + < path.size(); j++) {
            long key = edge(path.get(j), path.get(j + 1));
            x.put(key, hl.segmentTree.joinValueWithDelta(x.get(key), delta));
          }
        else {
          int res1 = hl.query(a, b);
          int res2 = hl.getNeutralValue();
          for (int j = 0; j + < path.size(); j++) {
            long key = edge(path.get(j), path.get(j + 1));
            res2 = hl.segmentTree.queryOperation(res2, x.get(key));
          }
          if (res1 != res2)
            throw new RuntimeException();
        }
      }
    }
    System.out.println("Test passed");
  }

  static long edge(int u, int v) {
    return ((longMath.min(u, v<< 16+ Math.max(u, v);
  }

  static boolean getPathFromAtoB(List<Integer>[] tree, int a, int b, int p, List<Integer> path) {
    path.add(a);
    if (a == b)
      return true;
    for (int u : tree[a])
      if (u != p && getPathFromAtoB(tree, u, b, a, path))
        return true;
    path.remove(path.size() 1);
    return false;
  }

  static List<Integer>[] getRandomTree(int n, Random rnd) {
    List<Integer>[] t = new List[n];
    for (int i = 0; i < n; i++)
      t[inew ArrayList<>();
    int[] p = new int[n];
    for (int i = 0, j; i < n; j = rnd.nextInt(i + 1), p[i= p[j], p[j= i, i++// random permutation
    for (int i = 1; i < n; i++) {
      int parent = p[rnd.nextInt(i)];
      t[parent].add(p[i]);
      t[p[i]].add(parent);
    }
    return t;
  }
}
Comments