Java 5ms iterative, following hint, O(klogn) time and space


  • 37
    Y

    Following the hint, I use a predecessor stack and successor stack. I do a logn traversal to initialize them until I reach the null node. Then I use the getPredecessor and getSuccessor method to pop k closest nodes and update the stacks.

    Time complexity is O(klogn), since k BST traversals are needed and each is bounded by O(logn) time. Note that it is not O(logn + k) which is the time complexity for k closest numbers in a linear array.

    Space complexity is O(klogn), since each traversal brings O(logn) new nodes to the stack.

    public class Solution {
        public List<Integer> closestKValues(TreeNode root, double target, int k) {
            List<Integer> result = new LinkedList<Integer>();
            // populate the predecessor and successor stacks 
            Stack<TreeNode> pred = new Stack<TreeNode>();
            Stack<TreeNode> succ = new Stack<TreeNode>();
            TreeNode curr = root;
            while (curr != null) {
                if (target <= curr.val) {
                    succ.push(curr);
                    curr = curr.left;
                } else {
                    pred.push(curr);
                    curr = curr.right;
                }
            }
            while (k > 0) {
                if (pred.empty() && succ.empty()) {
                    break; 
                } else if (pred.empty()) {
                    result.add( getSuccessor(succ) );
                } else if (succ.empty()) {
                    result.add( getPredecessor(pred) );
                } else if (Math.abs(target - pred.peek().val) < Math.abs(target - succ.peek().val)) {
                    result.add( getPredecessor(pred) );                    
                } else {
                    result.add( getSuccessor(succ) );
                }
                k--;
            }
            return result;
         }
    
        private int getPredecessor(Stack<TreeNode> st) {
            TreeNode popped = st.pop();
            TreeNode p = popped.left;
            while (p != null) {
                st.push(p);
                p = p.right;
            }
            return popped.val;
        }
    
        private int getSuccessor(Stack<TreeNode> st) {
            TreeNode popped = st.pop();
            TreeNode p = popped.right;
            while (p != null) {
                st.push(p);
                p = p.left;
            }
            return popped.val;
        }
    }

  • 0

    I think this is O(log n + k) because you don't do a full traversal k times. You only go as far as it is necessary to find the next element.


  • 0

    On the second thought, it looks like it is O(k log n) worst case and O(log n + k) best case. Because these traversals may or may not involve O(log n) loops. If we're lucky and these while loops in the helper functions terminate quickly, it'll be O(log n + k). If they run through the whole tree, that would be (k log n).


  • -1
    T

    The thinking is really awesome. However, will the code pass the test where target can be find in a non-leaf node and k = 1.
    Say:

    Tree:
           4
         /   \
        2     10
             /  \
            8    11
           / \
          7   9
    

    target = 10.00
    k = 1;


  • 1
    C

    @SergeyTachenov The whole tree would only be traversed at most once during the two while loop in the two get functions. The worst case will be O(n + log(n)).


  • 0
    A

    @table It will work since 10 will be at the top of succ


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.