Leetcode 2846. Minimum Edge Weight Equilibrium Queries in a Tree

[Woodyiiiiiii]

2846. Minimum Edge Weight Equilibrium Queries in a Tree

2846. Minimum Edge Weight Equilibrium Queries in a Tree

类似题目：
1483. Kth Ancestor of a Tree Node

LCA模版
【模板讲解】树上倍增算法（以及最近公共祖先）
LCA 模板（Python/Java/C++/Go）

class Solution {
    public int[] minOperationsQueries(int n, int[][] edges, int[][] queries) {
        int queriesLen = queries.length;
        int[] ans = new int[queriesLen];
        Map<Integer, Set<int[]>> g = new HashMap<>();
        for (int[] edge : edges) {
            int u = edge[0], v = edge[1], w = edge[2];
            g.putIfAbsent(u, new HashSet<>());
            g.putIfAbsent(v, new HashSet<>());
            g.get(u).add(new int[]{v, w});
            g.get(v).add(new int[]{u, w});
        }

        int m = 32 - Integer.numberOfLeadingZeros(n);
        int[][] pa = new int[n][m];
        for (int i = 0; i < n; i++) {
            Arrays.fill(pa[i], -1);
        }
        int[][][] cnt = new int[n][m][26];
        int[] depth = new int[n];
        dfs(0, -1, g, pa, cnt, depth);

        // init
        for (int j = 0; j < m - 1; j++) {
            for (int i = 0; i < n; i++) {
                int fa = pa[i][j];
                if (fa != -1) {
                    pa[i][j + 1] = pa[fa][j];
                    for (int k = 0; k < 26; k++) {
                        cnt[i][j + 1][k] += cnt[i][j][k] + cnt[fa][j][k];
                    }
                }
            }
        }

        // lca
        for (int i = 0; i < queriesLen; i++) {
            int x = queries[i][0], y = queries[i][1];
            if (depth[x] > depth[y]) {
                int tmp = x;
                x = y;
                y = tmp;
            }
            int pathLen = depth[y] + depth[x];
            int[] cntDiff = new int[26];

            // put x and y at same level
            for (int j = depth[y] - depth[x]; j > 0; j &= (j - 1)) {
                int k = Integer.numberOfTrailingZeros(j);
                for (int l = 0; l < 26; l++) {
                    cntDiff[l] += cnt[y][k][l];
                }
                y = pa[y][k];
            }

            if (x != y) {
                for (int j = m - 1; j >= 0; j--) {
                    int px = pa[x][j], py = pa[y][j];
                    if (px != py) {
                        for (int k = 0; k < 26; k++) {
                            cntDiff[k] += cnt[x][j][k] + cnt[y][j][k];
                        }
                        x = px;
                        y = py;
                    }
                }
                for (int k = 0; k < 26; k++) {
                    cntDiff[k] += cnt[x][0][k] + cnt[y][0][k];
                }
                x = pa[x][0];
            }

            int lca = x;
            pathLen -= 2 * depth[lca];
            int max = 0;
            for (int j = 0; j < 26; j++) {
                max = Math.max(max, cntDiff[j]);
            }
            ans[i] = pathLen - max;
        }

        return ans;
    }

    private void dfs(int u, int p, Map<Integer, Set<int[]>> g, int[][] pa, int[][][] cnt, int[] depth) {
        pa[u][0] = p;
        for (int[] edge : g.getOrDefault(u, new HashSet<>())) {
            int v = edge[0], w = edge[1];
            if (v == p) continue;
            depth[v] = depth[u] + 1;
            cnt[v][0][w - 1]++;
            dfs(v, u, g, pa, cnt, depth);
        }
    }
}