LCA 欧拉+ST表 Java实现

2023-03-25  本文已影响0人  kaiker
import java.util.ArrayList;
import java.util.List;

class LCA {
    private int[] euler; // 树的遍历结果
    private int[] depths; // 节点深度
    private int[] first; // 节点在遍历结果中首次出现的位置
    private int[] logTable; // log2表
    private int[][] st; // ST表
    private int idx; // 当前节点在遍历结果中的位置

    public LCA(TreeNode root, int nodeNum) {
        euler = new int[2 * nodeNum - 1];
        depths = new int[2 * nodeNum - 1];
        first = new int[nodeNum];
        logTable = new int[euler.length + 1];
        idx = 0;
        dfs(root, 0);
        buildLogTable();
        buildSTTable();
    }

    // 欧拉遍历
    private void dfs(TreeNode node, int depth) {
        first[node.val] = idx;
        euler[idx] = node.val;
        depths[idx++] = depth;
        for (TreeNode child : node.children) {
            dfs(child, depth + 1);
            euler[idx] = node.val;
            depths[idx++] = depth;
        }
    }

    private void buildLogTable() {
        for (int i = 2; i <= euler.length; i++) {
            logTable[i] = logTable[i / 2] + 1;
        }
    }

    private void buildSTTable() {
        int n = euler.length;
        int m = logTable[n] + 1;
        st = new int[n][m];
        for (int i = 0; i < n; i++) {
            st[i][0] = i;
        }
        for (int j = 1; j < m; j++) {
            for (int i = 0; i + (1 << j) <= n; i++) {
                int x = st[i][j - 1];
                int y = st[i + (1 << (j - 1))][j - 1];
                st[i][j] = depths[x] < depths[y] ? x : y;
            }
        }
    }

    private int querySTTable(int left, int right) {
        int k = logTable[right - left + 1];
        int x = st[left][k];
        int y = st[right - (1 << k) + 1][k];
        return depths[x] < depths[y] ? euler[x] : euler[y];
    }

    public int findLCA(TreeNode p, TreeNode q) {
        int left = first[p.val];
        int right = first[q.val];
        if (left > right) {
            int tmp = left;
            left = right;
            right = tmp;
        }
        return querySTTable(left, right);
    }

    public static void main(String[] args) {
        TreeNode node5 = new TreeNode(5);
        TreeNode node4 = new TreeNode(4);
        TreeNode node3 = new TreeNode(3);
        TreeNode node2 = new TreeNode(2);
        node2.children.add(node5);
        TreeNode node1 = new TreeNode(1);
        node1.children.add(node3);
        node1.children.add(node4);
        TreeNode node0 = new TreeNode(0);
        node0.children.add(node1);
        node0.children.add(node2);
        LCA lca = new LCA(node0, 6);
        System.out.println(lca.findLCA(node3, node4));
    }
}


class TreeNode {
    int val;
    List<TreeNode> children;
    // constructor
    public TreeNode(int val) {
        this.val = val;
        this.children = new ArrayList<>();
    }
}
上一篇下一篇

猜你喜欢

热点阅读