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


  • 32
    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).


  • 0
    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)).


Log in to reply
 

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