Skip to content

1483. Kth Ancestor of a Tree Node 👍

  • Time: $O(n\log n)$
  • Space: $O(n\log n)$
 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
class TreeAncestor {
 public:
  TreeAncestor(int n, vector<int>& parent)
      : maxLevel(32 - __builtin_clz(n)), dp(n, vector<int>(maxLevel)) {
    // Node i's 2^0 ancestor is its direct parent.
    for (int i = 0; i < n; ++i)
      dp[i][0] = parent[i];

    for (int j = 1; j < maxLevel; ++j)
      for (int i = 0; i < n; ++i)
        if (dp[i][j - 1] == -1)  // There's no such ancestor.
          dp[i][j] = -1;
        else  // A(i, 2^j) = A(A(i, 2^{j - 1}), 2^{j - 1})
          dp[i][j] = dp[dp[i][j - 1]][j - 1];
  }

  int getKthAncestor(int node, int k) {
    for (int j = 0; j < maxLevel && node != -1; ++j)
      if (k >> j & 1)
        node = dp[node][j];
    return node;
  }

 private:
  const int maxLevel;
  vector<vector<int>> dp;  // dp[i][j] := node i's 2^j-th ancestor
};

 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
class TreeAncestor {
  public TreeAncestor(int n, int[] parent) {
    this.maxLevel = 32 - Integer.numberOfLeadingZeros(n);
    this.dp = new int[n][maxLevel];

    // Node i's 2^0 ancestor is its direct parent.
    for (int i = 0; i < n; ++i)
      dp[i][0] = parent[i];

    for (int j = 1; j < maxLevel; ++j)
      for (int i = 0; i < n; ++i)
        if (dp[i][j - 1] == -1) // There's no such ancestor.
          dp[i][j] = -1;
        else // A(i, 2^j) = A(A(i, 2^{j - 1}), 2^{j - 1})
          dp[i][j] = dp[dp[i][j - 1]][j - 1];
  }

  public int getKthAncestor(int node, int k) {
    for (int j = 0; j < maxLevel && node != -1; ++j)
      if ((k >> j & 1) == 1)
        node = dp[node][j];
    return node;
  }

  private final int maxLevel;
  private int[][] dp; // dp[i][j] := node i's 2^j-th ancestor
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class TreeAncestor:
  def __init__(self, n: int, parent: list[int]):
    self.maxLevel = n.bit_length()
    # dp[i][j] := node i's 2^j-th ancestor
    self.dp = [[0] * self.maxLevel for _ in range(n)]

    # Node i's 2^0 ancestor is its direct parent
    for i in range(n):
      self.dp[i][0] = parent[i]

    for j in range(1, self.maxLevel):
      for i in range(n):
        if self.dp[i][j - 1] == -1:  # There's no such ancestor
          self.dp[i][j] = -1
        else:  # A(i, 2^j) = A(A(i, 2^{j - 1}), 2^{j - 1})
          self.dp[i][j] = self.dp[self.dp[i][j - 1]][j - 1]

  def getKthAncestor(self, node: int, k: int) -> int:
    for j in range(self.maxLevel):
      if node == -1:
        break
      if k >> j & 1:
        node = self.dp[node][j]
    return node