[LintCode] Problem 901 - Closest Binary Search Tree Value II

Posted on | In ,

Given a non-empty binary search tree and a target value, find k values in the BST that are closest to the target.

Note

  • Given target value is a floating point.
  • You may assume k is always valid, that is: k ≤ total nodes.
  • You are guaranteed to have only one unique set of k values in the BST that are closest to the target.

Example

No.1

Input:
{1}
0.000000
1

Output:
[1]

Explanation：
Binary tree {1}, denote the following structure:

1
1

No.2

Input:
{3,1,4,#,2}
0.275000
2

Output:
[1,2]

Explanation：
Binary tree {3,1,4,#,2}, denote the following structure:

1
2
3
4
5
  3
 /  \
1    4
 \
  2

Challenge

Assume that the BST is balanced, could you solve it in less than O(n) runtime (where n = total nodes)?

Code

1
2
3
4
5
6
public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int x) { val = x; }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
private Stack<TreeNode> predecessor = new Stack<>();
private Stack<TreeNode> successor = new Stack<>();

public List<Integer> closestKValues(TreeNode root, double target, int k) {
    List<Integer> result = new ArrayList<>();
    TreeNode node;
    getStack(root, target);

    while (k-- > 0) {
        if (successor.isEmpty() || !predecessor.isEmpty()
            && target - predecessor.peek().val < successor.peek().val - target) {
            node = predecessor.pop();
            result.add(node.val);
            getPredecessor(node);
        }
        else {
            node = successor.pop();
            result.add(node.val);
            getSuccessor(node);
        }
    }

    return result;
}

private void getStack(TreeNode root, double target) {
    while (root != null) {
        if (root.val < target) {
            predecessor.push(root);
            root = root.right;
        }
        else {
            successor.push(root);
            root = root.left;
        }
    }
}

private void getPredecessor(TreeNode root) {
    if (root.left == null)
        return;

    root = root.left;

    while (root != null) {
        predecessor.push(root);
        root = root.right;
    }
}

private void getSuccessor(TreeNode root) {
    if (root.right == null)
        return;

    root = root.right;

    while (root != null) {
        successor.push(root);
        root = root.left;
    }
}