272. Closest Binary Search Tree Value II

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:

Input: root = [4,2,5,1,3], target = 3.714286, and k = 2

    4
   / \
  2   5
 / \
1   3

Output: [4,3]

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

Rust Solution

struct Solution;
use rustgym_util::*;

trait Inorder {
    fn inorder(&self, values: &mut Vec<i32>);
}

impl Inorder for TreeLink {
    fn inorder(&self, values: &mut Vec<i32>) {
        if let Some(node) = self {
            let node = node.borrow();
            node.left.inorder(values);
            values.push(node.val);
            node.right.inorder(values);
        }
    }
}

impl Solution {
    fn closest_k_values(root: TreeLink, target: f64, k: i32) -> Vec<i32> {
        let mut values = vec![];
        root.inorder(&mut values);
        let n = values.len();
        let k = k as usize;
        let mut start = 0;
        let mut end = n;
        while end - start > k {
            if (values[start] as f64 - target).abs() > (values[end - 1] as f64 - target).abs() {
                start += 1;
            } else {
                end -= 1;
            }
        }
        values[start..end].to_vec()
    }
}

#[test]
fn test() {
    let root = tree!(4, tree!(2, tree!(1), tree!(3)), tree!(5));
    let target = 3.714286;
    let k = 2;
    let res = vec![3, 4];
    assert_eq!(Solution::closest_k_values(root, target, k), res);
}

Having problems with this solution? Click here to submit an issue on github.