给你一棵树,树上有 n 个节点,按从 0 到 n-1 编号。树以父节点数组的形式给出,其中 parent[i] 是节点 i 的父节点。树的根节点是编号为 0 的节点。
请你设计并实现 getKthAncestor(int node, int k) 函数,函数返回节点 node 的第 k 个祖先节点。如果不存在这样的祖先节点,返回 -1 。
树节点的第 k 个祖先节点是从该节点到根节点路径上的第 k 个节点。
示例:
输入: ["TreeAncestor","getKthAncestor","getKthAncestor","getKthAncestor"] [[7,[-1,0,0,1,1,2,2]],[3,1],[5,2],[6,3]]
输出: [null,1,0,-1]
解释: TreeAncestor treeAncestor = new TreeAncestor(7, [-1, 0, 0, 1, 1, 2, 2]);
treeAncestor.getKthAncestor(3, 1); // 返回 1 ,它是 3 的父节点 treeAncestor.getKthAncestor(5, 2); // 返回 0 ,它是 5 的祖父节点 treeAncestor.getKthAncestor(6, 3); // 返回 -1 因为不存在满足要求的祖先节点
提示:
1 <= k <= n <= 5*10^4 parent[0] == -1 表示编号为 0 的节点是根节点。 对于所有的 0 < i < n ,0 <= parent[i] < n 总成立 0 <= node < n 至多查询 5*10^4 次
解题思路:dp[node][j] 存储的是 node 节点距离为 2^j 的祖先是谁。
根据定义,dp[node][0] 就是 parent[node],即 node 的距离为 1 的祖先是 parent[node]。
状态转移是: dp[node][j] = dp[dp[node][j - 1]][j - 1]。
意思是:要想找到 node 的距离 2^j 的祖先,先找到 node 的距离 2^(j - 1) 的祖先,然后,再找这个祖先的距离 2^(j - 1) 的祖先。两步得到 node 的距离为 2^j 的祖先。
所以,我们要找到每一个 node 的距离为 1, 2, 4, 8, 16, 32, ... 的祖先,直到达到树的最大的高度。树的最大的高度是 logn 级别的。
这样做,状态总数是 O(nlogn),可以使用 O(nlogn) 的时间做预处理。
之后,根据预处理的结果,可以在 O(logn) 的时间里完成每次查询:对于每一个查询 k,把 k 拆解成二进制表示,然后根据二进制表示中 1 的位置,累计向上查询。
class TreeAncestor {
public:
vector<vector<int>>dp;
TreeAncestor(int n, vector<int>& parent):dp(n) {
for(int i=0;i<parent.size();i++)
{
dp[i].push_back(parent[i]);
}
for(int j=1;;j++)
{
bool flag=1;
for(int i=0;i<n;i++)
{
int t=dp[i][j-1]==-1?-1:dp[dp[i][j-1]][j-1];
dp[i].push_back(t);
if(t!=-1)flag=0;
}
if(flag)break;
}
}
int getpos(int x)
{
int k=x,ans=0;
while(k)
{
if(k&1)return ans;
else
{
ans++;
k>>=1;
}
}
return k;
}
int getKthAncestor(int node, int k) {
//cout<<node<<" "<<k<<" "<<getpos(k)<<" "<<dp[node][getpos(k)]<<" "<<k-(1<<getpos(k))<<endl;
if(node==-1||k==0)return node;
if(getpos(k)>dp[node].size()-1)return -1;
return getKthAncestor(dp[node][getpos(k)],k-(1<<getpos(k)));
}
};
/**
* Your TreeAncestor object will be instantiated and called as such:
* TreeAncestor* obj = new TreeAncestor(n, parent);
* int param_1 = obj->getKthAncestor(node,k);
*/
我的博客即将同步至腾讯云+社区,邀请大家一同入驻:https://cloud.tencent.com/developer/support-plan?invite_code=1mn7k3k6lbl1o