Leetcode 1483. 樹節點的第 K 個祖先(DP,倍增)

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

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章