LeetCode: Count Complete Tree Nodes

以下の問題が面白かったのでメモ。

LeetCode - Count Complete Tree Nodes [Mediam] https://leetcode.com/problems/count-complete-tree-nodes/

以下のような再帰の解法で，別にACはされるのだが，Complete Tree Nodeの条件はいかせていない。

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    public int countNodes(TreeNode root) {
        if (root == null) {
            return 0;
        }
        return 1 + countNodes(root.left) + countNodes(root.right);
    }
}

とりあえず答え合わせ。なるほどなーという感じ。答えのアイデアはパクりつつ，できるだけ自分で書くようにしてみたら，以下のような感じで，解答と同じ計算量にできた。

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    
    public boolean exists(TreeNode root, int depth, int index, int max) {
        while (root != null) {
            int half = max / 2;
            if (index < half) {
                root = root.left;
            } else {
                root = root.right;
                index -= half;
            }
            max = half;
            depth--;
        }
        return depth == 0;
    }
    
    public int countNodes(TreeNode root) {
        int d = depth(root);
        if (d <= 1) {
            return d;
        }
        
        int leaves = (int) Math.pow(2, d - 1);
        
        int l = 0;
        int h = leaves - 1;
        while (l < h) {
            int m = l + (h - l + 1) / 2;
            if (exists(root, d, m, leaves)) {
                l = m;
            } else {
                h = m - 1;
            }
        }
        return (int) leaves - 1 + (l + 1);
    }
    
    public int depth(TreeNode n) {
        int d = 0;
        while (n != null) {
            n = n.left;
            d++;
        }
        return d;
    }
}

ざっくり解説すると，完全木では葉の深さは高々1の差しかなく，左端の葉が最も深いという特徴がある。

そのため，深さの最大値dとするとまではすべてのノードがあるといえる。深さd-1までのノード数は等比数列の和の公式からと求まる。 そして，深さのノードがいくつあるかを二分探索で求める。 深さのノード数の最大値はなので，二部探索の計算量は，になる。探索一回あたりの計算量はなので，か。