You are given a tree with n
nodes numbered from 0
to n-1
in the form of a parent array where parent[i]
is the parent of node i
. The root of the tree is node 0
.
Implement the function getKthAncestor
(int node, int k)
to return the k
-th ancestor of the given node
. If there is no such ancestor, return -1
.
The k-th ancestor of a tree node is the k
-th node in the path from that node to the root.
Example:
Input: ["TreeAncestor","getKthAncestor","getKthAncestor","getKthAncestor"] [[7,[-1,0,0,1,1,2,2]],[3,1],[5,2],[6,3]] Output:
[null,1,0,-1]
Explanation: TreeAncestor treeAncestor = new TreeAncestor(7, [-1, 0, 0, 1, 1, 2, 2]); treeAncestor.getKthAncestor(3, 1); // returns 1 which is the parent of 3 treeAncestor.getKthAncestor(5, 2); // returns 0 which is the grandparent of 5 treeAncestor.getKthAncestor(6, 3); // returns -1 because there is no such ancestor
Constraints:
1 <= k <= n <= 5*10^4
parent[0] == -1
indicating that0
is the root node.0 <= parent[i] < n
for all0 < i < n
0 <= node < n
- There will be at most
5*10^4
queries.
Solution: LogN ancestors
- Build the tree from parent array
- Traverse the tree
- For each node stores up to logn ancestros, 2^0-th, 2^1-th, 2^2-th, …
When k comes in, each node take the highest bit h out, and query its 2^h’s ancestors with k <- (k – 2^h). There will be at most logk recursive query. When it ends? k == 0, we found the ancestors which is the current node. Or node == 0 and k > 0, we already at root which doesn’t have any ancestors so return -1.
Time complexity:
Construction: O(nlogn)
Query: O(logk)
Space complexity:
O(nlogn)
C++
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 |
// Author: Huahua class TreeAncestor { public: TreeAncestor(int n, vector<int>& parent): g_(n), a_(n) { for (int i = 1; i < n; ++i) g_[parent[i]].push_back(i); vector<int> path; dfs(0, path); } int getKthAncestor(int node, int k) { if (k == 0) return node; if (node == 0 && k > 0) return -1; int l = min(31ul - __builtin_clz(k), a_[node].size() - 1); return getKthAncestor(a_[node][l], k - (1 << l)); } private: vector<vector<int>> g_, a_; void dfs(int node, vector<int>& path) { for (int i = 1; i <= path.size(); i *= 2) a_[node].push_back(path[path.size() - i]); path.push_back(node); for (int c : g_[node]) dfs(c, path); path.pop_back(); } }; |
DP method
C++
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
// Author: Huahua // Author: Huahua class TreeAncestor { public: TreeAncestor(int n, vector<int>& parent): max_level_(32 - __builtin_clz(n)), dp_(max_level_, vector<int>(n)) { dp_[0] = parent; for (int i = 1; i < max_level_; ++i) for (int j = 0; j < n; ++j) dp_[i][j] = dp_[i - 1][j] == -1 ? -1 : dp_[i - 1][dp_[i - 1][j]]; } int getKthAncestor(int node, int k) { for (int i = 0; i < max_level_ && node != -1; ++i) if (k & (1 << i)) node = dp_[i][node]; return node; } private: const int max_level_; vector<vector<int>> dp_; }; |
Solution 2: Binary Search
credit: Ziwu Zhou
Construction: O(n)
Traverse the tree in post order, for each node record its depth and id (visiting order).
For each depth, store all the nodes and their ids.
Query: O(logn)
Get the depth and id of the node, if k > d, return -1.
Use binary search to find the first node at depth[d – k] that has a id greater than the query’s one That node is the k-th ancestor of the node.
C++
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 |
// Author: Huahua class TreeAncestor { public: TreeAncestor(int n, vector<int>& parent): g_(n), nodes_(n), levels_(n) { for (int i = 1; i < n; ++i) g_[parent[i]].push_back(i); int id = 0; dfs(0, 0, id); } int getKthAncestor(int node, int k) { const auto [d, id] = nodes_[node]; if (k > d) return -1; const auto& ns = levels_[d - k]; return upper_bound(begin(ns), end(ns), pair{id, -1})->second; } private: void dfs(int node, int depth, int& id) { for (int c : g_[node]) dfs(c, depth + 1, ++id); levels_[depth].emplace_back(id, node); nodes_[node] = {depth, id}; } vector<vector<int>> g_; // p -> {{child}} vector<vector<pair<int, int>>> levels_; // depth -> {{id, node}} vector<pair<int, int>> nodes_; // node -> {depth, id} }; |
请尊重作者的劳动成果,转载请注明出处!花花保留对文章/视频的所有权利。
如果您喜欢这篇文章/视频,欢迎您捐赠花花。
If you like my articles / videos, donations are welcome.
Be First to Comment