算法第十七週作業01

Description

Given a binary search tree, write a function kthSmallest to find the kth smallest element in it.

Solution

1. 獲取左子樹的節點數number
   如果左子樹節點數剛好爲k-1，說明該節點爲第k個，否則向左子樹找或者向右子樹找。

Code

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public int kthSmallest(TreeNode root, int k) {
        int number = number(root.left);

        if (number == k - 1) {
            // 如果左子樹節點數剛好爲k-1，說明該節點爲第k個
            return root.val;
        } else if (number >= k) {
            // 第k個在左子樹
            return kthSmallest(root.left, k);
        } else {
            // 第k個在右子樹
            return kthSmallest(root.right, k - 1 - number);
        }
    }

    // 獲取該節點的所有葉子數
    public int number(TreeNode node) {
        if (node == null) {
            return 0;
        } else {
            return number(node.left) + number(node.right) + 1;
        }
    }
}