[1483. 樹節點的第 K 個祖先] 【路徑】

fishcanfly發表於2024-04-06

思路很簡單:


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

class TreeAncestor {
    List<Integer>[] children;
    List<List<Integer>> paths = new ArrayList<>();
    int[][] pathIdMap;

    public static void main(String[] args) {
        TreeAncestor ancestor = new TreeAncestor(5, new int[]{
                -1, 0, 0, 0, 3
        });
        int ans = ancestor.getKthAncestor(3, 2);
        System.out.println(ans);
    }

    public TreeAncestor(int n, int[] parent) {
        children = new ArrayList[n];
        for (int i = 0; i < n; i++) {
            children[i] = new ArrayList<>();
        }
        for (int i = 0; i < n; i++) {
            if (i == 0) {
                continue;
            }
            children[parent[i]].add(i);
        }

        int pathId = 0;
        pathIdMap = new int[n + 1][2];
        for (int i = 0; i < n; i++) {
            if (!children[i].isEmpty()) {
                continue;
            }
            List<Integer> path = new ArrayList<>();
            int val = i;
            int num = 0;
            while (val != -1) {
                path.add(val);
                pathIdMap[val][0] = pathId;
                pathIdMap[val][1] = num++;
                val = parent[val];
            }
            paths.add(path);
            pathId++;
        }
    }

    public int getKthAncestor(int node, int k) {
        int[] pathId = pathIdMap[node];
        List<Integer> path = paths.get(pathId[0]);
        int size = path.size();
        int offset = pathId[1];
        if (size <= offset + k) {
            return -1;
        }
        return path.get(offset + k);
    }
}

/**
 * Your TreeAncestor object will be instantiated and called as such:
 * TreeAncestor obj = new TreeAncestor(n, parent);
 * int param_1 = obj.getKthAncestor(node,k);
 */

相關文章