算法練習 Leetcode 222 Count Complete Tree Nodes 解法

feellikelearning February 10, 2020 1 Comment

題目如下：

Definition of a complete binary tree from Wikipedia:
In a complete binary tree every level, except possibly the last, is completely filled, and all nodes in the last level are as far left as possible. It can have between 1 and 2^h nodes inclusive at the last level h.

Example:

Input: 
    1
   / \
  2   3
 / \  /
4  5 6

Output: 6

最簡單的做法就是直接遍歷一邊，不管是DFS還是BFS，時間複雜度都是O(N)。但是這樣就沒利用到complete binary tree的條件。優化的解法如果沒練過的話其實不太好想。這題要的其實是找到最後一層有幾個node。而每層的node的index就是在[2^h, 2^(h+1)-1]區間之內。如果有個function可以告訴你某個index在不在tree里，那麼可以對[2^h, 2^(h+1)-1]做binary search，找到其中index最大並且在tree里的node。Binary tree的最後一層大概有N/2個node（每加一層，node數量翻倍！），那麼需要做O(lgN)次search。

至於這個告訴某個node在不在tree里的函數可以這樣寫：寫幾個例子觀察一下，如果一個node的index是n, 每個node的parent的index都是floor(n/2)。那麼我們就能把從root到達目標node的path找出來，比如如果目標node是5，那麼path就是[5, 2, 1]，入股目標是8，那麼path就是[8, 4, 2, 1]。然後從root，也就是1號node出發，如果最後我們按照這個地圖能走到目標，那麼這個node就在tree裏面，否則就不在。這個function需要O(lgN)時間和O(lgN)空間。

最後總的時間複雜度這樣計算，O(lgN)次search，每次 O(lgN)時間，總時間O((lgN)^2)。空間O(lgN)用來保存path。

代碼如下

C++

class Solution {
public:
    int countNodes(TreeNode* root) {
        if (root == NULL) {
            return 0;
        }
        int h = findHeight(root);
        int i_min = pow(2, h);
        int i_max = pow(2, h + 1) - 1;
        int ans = i_min;
        while(i_min <= i_max) {
            int m = i_min + (i_max - i_min) / 2;
            if (inTree(root, m)) {
                ans = m;
                i_min = m + 1;
            } else{
                i_max = m - 1;
            }
        }   
        return ans;
        
    }
    
    int findHeight(TreeNode* root) {
        int h = 0;
        TreeNode* cur = root;
        while (cur != NULL) {
            cur = cur->left;
            h += 1;
        }
        return h - 1;
    }
    
    bool inTree(TreeNode* root, int i) {
        if (root == NULL || i < 0) {
            return false;
        }
        vector<int> path;
        int cur = i;
        while (cur != 0) {
            path.push_back(cur);
            cur /= 2;
        }
        TreeNode* curNode = root;
        for (int j = path.size() - 2; j >= 0; j--) {
            if (path[j] == path[j + 1] * 2) {
                curNode = curNode->left;
            } else {
                curNode = curNode->right;
            }
            if (curNode == NULL) {
                return false;
            }
        }
        return true;
    }
};

Java

class Solution {
    public int countNodes(TreeNode root) {
        if (root == null) {
            return 0;
        }
        int h = findHeight(root);
        int i_min = (int)Math.pow(2, h);
        int i_max = (int)(Math.pow(2, h + 1) - 1);
        int ans = i_min;
        while (i_min <= i_max) {
            int m = i_min + (i_max - i_min) / 2;
            if (inTree(root, m)) {
                ans = m;
                i_min = m + 1;
            } else{
                i_max = m - 1;
            }
        }   
        return ans;
        
    }
    
    public int findHeight(TreeNode root) {
        int h = 0;
        TreeNode cur = root;
        while (cur != null) {
            cur = cur.left;
            h += 1;
        }
        return h - 1;
    }
    
    public boolean inTree(TreeNode root, int i) {
        if (root == null || i < 0) {
            return false;
        }
        List<Integer> path = new ArrayList<>();
        int cur = i;
        while (cur != 0) {
            path.add(cur);
            cur /= 2;
        }
        TreeNode curNode = root;
        for (int j = path.size() - 2; j >= 0; j--) {
            if (path.get(j) == path.get(j + 1) * 2) {
                curNode = curNode.left;
            } else {
                curNode = curNode.right;
            }
            if (curNode == null) {
                return false;
            }
        }
        return true;
    }
}