ClosestBinarySearchTreeValue II: done

Author: rampatra

src/main/java/com/leetcode/trees/ClosestBinarySearchTreeValueII.java

@@ -1,27 +1,121 @@
 package com.leetcode.trees;
 
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Queue;
+import java.util.Stack;
+
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 /**
  * Level: Hard
  * Problem Link: https://leetcode.com/problems/closest-binary-search-tree-value-ii/
  * Problem Description:
+ * 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)?
  *
  * @author rampatra
  * @since 2019-07-31
  */
 public class ClosestBinarySearchTreeValueII {
 
+
+    /**
+     * The idea is simple. We do the inorder traversal and keep the values less than or equal to target in a stack and
+     * the values greater than target in a queue. And finally, we compare values from both stack and queue and take
+     * whichever is the closest to target value each time.
+     *
+     * Note: We can optimize it even further in terms of space. We can get rid of the stack and queue and just fill up
+     * the result list in the recursive inOrder call. Once the result list is of size k, we can compare and remove the
+     * farthest value and insert the closer value. See {@link ClosestBinarySearchTreeValueII#closestKValuesOptimized(TreeNode, double, int)}.
+     *
+     * @param root
+     * @param target
+     * @param k
+     * @return
+     */
+    public static List<Integer> closestKValues(TreeNode root, double target, int k) {
+        int count = 0;
+        List<Integer> closestKValues = new LinkedList<>();
+
+        Stack<Integer> predecessors = new Stack<>();
+        Queue<Integer> successors = new LinkedList<>();
+        inOrder(root, predecessors, successors, target, k);
+
+        while (count < k) {
+            if (predecessors.empty()) {
+                closestKValues.add(successors.poll());
+            } else if (successors.isEmpty()) {
+                closestKValues.add(predecessors.pop());
+            } else if (Math.abs(target - predecessors.peek()) < Math.abs(target - successors.peek())) {
+                closestKValues.add(predecessors.pop());
+            } else {
+                closestKValues.add(successors.poll());
+            }
+            count++;
+        }
+
+        return closestKValues;
+    }
+
+    private static void inOrder(TreeNode root, Stack<Integer> predecessors, Queue<Integer> successors, double target, int k) {
+        if (root == null || successors.size() == k) return;
+        inOrder(root.left, predecessors, successors, target, k);
+        if (root.val <= target) {
+            predecessors.add(root.val);
+        } else {
+            successors.add(root.val);
+        }
+        inOrder(root.right, predecessors, successors, target, k);
+    }
+
+
     /**
-     * @param node
-     * @param parentNode
-     * @param val
-     * @param diff
+     * This approach is similar to the above one but it doesn't use stack or queue.
+     *
+     * @param root
+     * @param target
+     * @param k
      * @return
      */
-    public static TreeNode findNodeWithClosestValue(TreeNode node, TreeNode parentNode, int val, int diff) {
-        return null;
+    public static List<Integer> closestKValuesOptimized(TreeNode root, double target, int k) {
+        LinkedList<Integer> closestKValues = new LinkedList<>();
+        inOrder(root, target, k, closestKValues);
+        return closestKValues;
+    }
+
+    private static void inOrder(TreeNode root, double target, int k, LinkedList<Integer> closestKValues) {
+        if (root == null) return;
+
+        inOrder(root.left, target, k, closestKValues);
+        if (closestKValues.size() == k) {
+            //if size k, add current and remove head if it's closer to target, otherwise return
+            if (Math.abs(target - root.val) < Math.abs(target - closestKValues.peekFirst()))
+                closestKValues.removeFirst();
+            else {
+                return;
+            }
+        }
+        closestKValues.add(root.val);
+        inOrder(root.right, target, k, closestKValues);
     }
 
     public static void main(String[] args) {
@@ -46,13 +140,8 @@ public static void main(String[] args) {
         root.left.left.right = new TreeNode(6);
         root.right.right = new TreeNode(20);
 
-        assertEquals(13, findNodeWithClosestValue(root, root, 15, Integer.MAX_VALUE).val);
-        assertEquals(13, findNodeWithClosestValue(root, root, 13, Integer.MAX_VALUE).val);
-        assertEquals(9, findNodeWithClosestValue(root, root, 9, Integer.MAX_VALUE).val);
-        assertEquals(2, findNodeWithClosestValue(root, root, 2, Integer.MAX_VALUE).val);
-        assertEquals(2, findNodeWithClosestValue(root, root, 1, Integer.MAX_VALUE).val);
-        assertEquals(6, findNodeWithClosestValue(root, root, 6, Integer.MAX_VALUE).val);
-        assertEquals(13, findNodeWithClosestValue(root, root, 11, Integer.MAX_VALUE).val); // tie b/w 9 and 13
+        assertEquals("[9, 8, 7, 6, 5]", closestKValues(root, 8.5, 5).toString());
+        assertEquals("[5, 6, 7, 8, 9]", closestKValuesOptimized(root, 8.5, 5).toString());
 
         /*
             BST looks like:
@@ -67,10 +156,11 @@ public static void main(String[] args) {
         root.left = new TreeNode(7);
         root.right = new TreeNode(13);
         root.left.left = new TreeNode(5);
-        root.left.right = new TreeNode(8);
         root.right.left = new TreeNode(13);
+        root.left.right = new TreeNode(8);
         root.right.right = new TreeNode(20);
 
-        assertEquals(13, findNodeWithClosestValue(root, root, 15, Integer.MAX_VALUE).val);
+        assertEquals("[13, 13, 9, 20, 8]", closestKValues(root, 14, 5).toString());
+        assertEquals("[8, 9, 13, 13, 20]", closestKValuesOptimized(root, 14, 5).toString());
     }
 }