AC clean Java solution using two stacks


  • 144

    The idea is to compare the predecessors and successors of the closest node to the target, we can use two stacks to track the predecessors and successors, then like what we do in merge sort, we compare and pick the closest one to the target and put it to the result list.

    As we know, inorder traversal gives us sorted predecessors, whereas reverse-inorder traversal gives us sorted successors.

    We can use iterative inorder traversal rather than recursion, but to keep the code clean, here is the recursion version.

    public List<Integer> closestKValues(TreeNode root, double target, int k) {
      List<Integer> res = new ArrayList<>();
    
      Stack<Integer> s1 = new Stack<>(); // predecessors
      Stack<Integer> s2 = new Stack<>(); // successors
    
      inorder(root, target, false, s1);
      inorder(root, target, true, s2);
      
      while (k-- > 0) {
        if (s1.isEmpty())
          res.add(s2.pop());
        else if (s2.isEmpty())
          res.add(s1.pop());
        else if (Math.abs(s1.peek() - target) < Math.abs(s2.peek() - target))
          res.add(s1.pop());
        else
          res.add(s2.pop());
      }
      
      return res;
    }
    
    // inorder traversal
    void inorder(TreeNode root, double target, boolean reverse, Stack<Integer> stack) {
      if (root == null) return;
    
      inorder(reverse ? root.right : root.left, target, reverse, stack);
      // early terminate, no need to traverse the whole tree
      if ((reverse && root.val <= target) || (!reverse && root.val > target)) return;
      // track the value of current node
      stack.push(root.val);
      inorder(reverse ? root.left : root.right, target, reverse, stack);
    }

  • 0
    C

    Nice clean code!


  • 1
    O

    This is O(n + k), right?


  • 1

    Yes, it's O(n + k) time complexity.


  • 0
    A

    Can I do inorder traverse first, and apply binary search? That'd be O(logn+k) time complexity, but with O(n) extra space.


  • 1

    A full inorder traversal takes O(n) time already, plus binary search and pick, total time would be O(n + logn + k).


  • 1
    K

    What does the followup question try to imply?
    Follow up:
    Assume that the BST is balanced, could you solve it in less than O(n) runtime (where n = total nodes)?


  • 1
    M

    If I understand correctly, it seems first call inorder to get all the number less then target, then calll reverse inorder to get all the number greater the target. So it need O(n), not O(logN)??
    So why not do one inoder and get total ordered array, and then using binary search find the nearest target and split the total array?


  • 5

    Hi, mengjunxie.us. I guess using your idea, we will need three O(n) scans: (1) the inorder traversal; (2) find the closest single value; (3) find all the k values. And the above code only needs to O(n) scans: (1) inorder traversal; (2) find all the k values, since all the nodes have been split into two halves at the same time during the traversal.


  • 3

    Great code! I really like the early stop :-)

    if ((reverse && root.val <= target) || (!reverse && root.val > target)) return;

  • 0
    C

    I love your code


  • 0
    W

    hmm, I have the same idea with you@xiemengjun.us.I think using stack implementation will avoid a binary search time maybe.


  • 0
    A

    Hi, Jianchao.li.fighter,
    You mentioned to mengjunxie.us that his proposed algo will take three O(N) scans, well I disagree with your reasoning.
    "find the closest single value;" -> Will take O(log n) by binary search as its sorted now.


  • 32
    Y

    your solution is O(n). And you have made simple things complex. For a O(n) solution, we don't need two stack. We also dont need two scan : in order and reverse in order. We just do a single in order transverse, and maintain a window size K.

    We should find a O(K) solution for this problem.


  • 2

    Hi yueyub, my solution is O(n + k) not O(n), when you say you can find a O(k) solution for this problem, I disagree. Let's say we have a BST with one million nodes, now to find the closest 2 nodes to the target, you say we can handle in O(k) time which is 2, how? I think that's impossible, we need to scan all the nodes in a single in order traversal, and that takes O(n) time already. Also, I am curious how you are going to "maintain" the window of size k using what data structure (array, deque, priority queue or something else).


  • 1
    Y

    sorry for the inaccurate statement of O(K). To be more precise, we should find a O(log(n)+K) solution. We use log(n) time to locate the "upwind" and "downwind" pointers for the target value. Then we scan the tree from the center to two sides. In this way, we can achieve the optimal solution, which I think is the key point of this problem and is also what interviewers looking for.


  • 1
    S

    Agree with yueyub. That's the correct solution. Using O(N) is too expensive and too easy for this problem since it is labeled as "Hard".


  • 0
    S

    Actually yueyub, do you have coded the solution for this problem? I have written mine but even it passes the all test cases, it seems a little complicated than I thought.


  • 0

    Anyone can share the better solution (with code) here? Thanks!


  • 0
    Y

    can we reverse the order of traversal, and then do an early stopping on the stack when each of the stack reaches size k?


Log in to reply
 

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