leetcode/Problem5456/Solution.java at master · gorthitk/leetcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

class TreeAncestor
{
    private final int[] parent;
    private final int[] depth;

    public TreeAncestor(int n, int[] parent)
    {
        this.parent = parent;
        this.depth = new int[n];

        for (int i = 1; i < parent.length; i++)
        {
            _findDepth(i);
        }
    }

    private int _findDepth(int node)
    {
        if (node == 0)
        {
            return 0;
        }

        if (depth[node] != 0)
        {
            return depth[node];
        }

        depth[node] = _findDepth(parent[node]) + 1;
        return depth[node];
    }

    public int getKthAncestor(int node, int k)
    {
        if (depth[node] < k)
        {
            return -1;
        }

        int ancestor = node;
        while (k > 0 && ancestor >= 0)
        {
            ancestor = parent[ancestor];
            k--;
        }

        return ancestor;
    }
}