一刷 05/2022
Version #1 PriorityQueue with BFS
Time O(Nlogk)
Space O(N)
/**
* Definition for a binary tree node.
* public class TreeNode {
* int val;
* TreeNode left;
* TreeNode right;
* TreeNode() {}
* TreeNode(int val) { this.val = val; }
* TreeNode(int val, TreeNode left, TreeNode right) {
* this.val = val;
* this.left = left;
* this.right = right;
* }
* }
*/
class Solution {
public List<Integer> closestKValues(TreeNode root, double target, int k) {
if (root == null) {
return new ArrayList<>();
}
Queue<TreeNode> que = new ArrayDeque<>();
Queue<Integer> maxHeap = new PriorityQueue<>((a, b) -> Math.abs(a - target) > Math.abs(b - target) ? -1 : 1);
que.offer(root);
while (!que.isEmpty()) {
TreeNode curr = que.poll();
maxHeap.offer(curr.val);
if (maxHeap.size() > k) {
maxHeap.remove();
}
if (curr.left != null) {
que.offer(curr.left);
}
if (curr.right != null) {
que.offer(curr.right);
}
}
return new ArrayList<>(maxHeap);
}
}
Version #2 In-order traversal with Deque
Time O(N)
Space O(N) - stack N, deque k
/**
* Definition for a binary tree node.
* public class TreeNode {
* int val;
* TreeNode left;
* TreeNode right;
* TreeNode() {}
* TreeNode(int val) { this.val = val; }
* TreeNode(int val, TreeNode left, TreeNode right) {
* this.val = val;
* this.left = left;
* this.right = right;
* }
* }
*/
class Solution {
public List<Integer> closestKValues(TreeNode root, double target, int k) {
// Use a ArrayDeque to store the candidates
// Do a in-order traversal
// Compare the current value difference with the first diff in the deque
// When the deque size grows to k:
// If current diff is smaller than the first diff, then remove the first element from the deque and replace with the current element
// If current diff is greater than the first diff, then we've already found the closest k element
if (root == null) {
return new ArrayList<>();
}
Deque<Integer> deque = new ArrayDeque<>();
Stack<TreeNode> stack = new Stack<>();
pushLeft(root, stack);
while (!stack.isEmpty()) {
TreeNode curr = stack.pop();
double diff = Math.abs(curr.val - target);
if (deque.size() == k) {
double firstDiff = Math.abs(deque.getFirst() - target);
if (diff > firstDiff) {
break;
}
deque.removeFirst();
}
deque.addLast(curr.val);
pushLeft(curr.right, stack);
}
return new ArrayList<>(deque);
}
private void pushLeft(TreeNode node, Stack<TreeNode> stack) {
while (node != null) {
stack.push(node);
node = node.left;
}
}
}
No comments:
Post a Comment