給你一棵樹,樹上有 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);
*/