Отбор на Чемпионат юга 2018. Задача F

package net.egork;

import net.egork.utils.io.InputReader;
import net.egork.utils.io.OutputWriter;

import java.util.ArrayList;
import java.util.List;

public class TaskF {
    int n;
    List<Integer>[] g;
    int[] size;
    int[] depth;
    int[] parent;

    public void solve(int testNumber, InputReader in, OutputWriter out) {
        n = in.readInt();
        g = new List[n];
        for (int i = 0; i < n; i++) {
            g[i] = new ArrayList<>();
        }
        for (int i = 0; i < n - 1; i++) {
            int u = in.readInt() - 1;
            int v = in.readInt() - 1;
            g[u].add(v);
            g[v].add(u);
        }
        size = new int[n];
        depth = new int[n];
        parent = new int[n];
        dfs(0);
        int wayCount = in.readInt();
        HLD hld = new HLD();
        int[][] ways = new int[wayCount][2];
        for (int i = 0; i < wayCount; i++) {
            ways[i][0] = in.readInt() - 1;
            ways[i][1] = in.readInt() - 1;
            int lca = hld.lca(ways[i][0], ways[i][1]);
            hld.add(lca, 1);
        }
        long answer = 0;
        for (int i = 0; i < wayCount; i++) {
            answer += hld.getSum(ways[i][0], ways[i][1]) - 1;
        }
        for (int i = 0; i < n; i++) {
            long sum = hld.getSum(i, i);
            answer -= sum * (sum - 1) / 2;
        }
        long result = wayCount * (wayCount - 1L) / 2 - answer;
        out.print(result);
    }

    void dfs(int v) {
        size[v] = 1;
        for (int to : g[v]) {
            g[to].remove((Integer) v);
            depth[to] = depth[v] + 1;
            parent[to] = v;
            dfs(to);
            size[v] += size[to];
        }
    }

    class Fenwick {
        int[] tree;

        Fenwick(int size) {
            tree = new int[size + 1];
        }

        void add(int at, int value) {
            for (at++; at < tree.length; at += at & -at) {
                tree[at] += value;
            }
        }

        int get(int at) {
            int s = 0;
            for (at++; at > 0; at -= at & -at) {
                s += tree[at];
            }
            return s;
        }

        int get(int l, int r) {
            return get(r) - get(l - 1);
        }
    }

    class HLD {

        List<List<Integer>> paths;
        int[] treeIndex;
        int[] indexInTree;
        Fenwick[] trees;

        HLD() {
            paths = new ArrayList<>();
            paths.add(new ArrayList<>());
            indexInTree = new int[n];
            treeIndex = new int[n];
            initTrees(0, 0);

            initTrees();
        }

        void initTrees(int root, int tree) {
            indexInTree[root] = paths.get(tree).size();
            paths.get(tree).add(root);
            treeIndex[root] = tree;
            for (int to : g[root]) {
                if (size[to] * 2 >= size[root]) {
                    initTrees(to, tree);
                } else {
                    paths.add(new ArrayList<>());
                    initTrees(to, paths.size() - 1);
                }
            }
        }

        int head(int path) {
            return paths.get(treeIndex[path]).get(0);
        }

        void initTrees() {
            trees = new Fenwick[paths.size()];
            for (int i = 0; i < trees.length; i++) {
                trees[i] = new Fenwick(paths.get(i).size());
            }
        }

        int getSum(int u, int v) {
            int s = 0;
            while (treeIndex[u] != treeIndex[v]) {
                if (depth[head(u)] > depth[head(v)]) {
                    int t = u;
                    u = v;
                    v = t;
                }
                s += trees[treeIndex[v]].get(indexInTree[v]);
                v = parent[head(v)];
            }
            s += trees[treeIndex[u]].get(
                    Math.min(indexInTree[u], indexInTree[v]),
                    Math.max(indexInTree[u], indexInTree[v])
            );
            return s;
        }

        int lca(int u, int v) {
            while (treeIndex[u] != treeIndex[v]) {
                if (depth[head(u)] > depth[head(v)]) {
                    int t = u;
                    u = v;
                    v = t;
                }
                v = parent[head(v)];
            }
            return depth[v] < depth[u] ? v : u;
        }

        void add(int at, int value) {
            Fenwick tree = trees[treeIndex[at]];
            tree.add(indexInTree[at], value);
        }
    }
}