BST中最接近K的数
问题描述
给定一棵BST和一个数K,找出BST中与K值最为接近的数。
解决思路
递归。
返回当前节点和目标值的差值、左子树节点和目标值的差值以及右子树节点和目标值的差值中,最小的那个所对应节点值。
程序
public class FindClosestInBST { public TreeNode findClosestNode(TreeNode root, int k) { if (root == null) { return null; } if (root.val == k) { return root; } TreeNode left = findClosestNode(root.left, k); TreeNode right = findClosestNode(root.right, k); int rootDiff = Math.abs(root.val - k); int leftDiff = left == null ? Integer.MAX_VALUE : Math .abs(left.val - k); int rightDiff = right == null ? Integer.MAX_VALUE : Math.abs(right.val - k); int min = Math.min(rootDiff, Math.min(leftDiff, rightDiff)); if (min == rootDiff) { return root; } else if (min == leftDiff) { return left; } return right; } }
Follow up
在BST中找与target最为接近的K个数。
解决思路
使用一个最小堆,存的是遍历元素与target的差值,比较的方法是比较差值的绝对值。
最后输出target和堆里面的差值之和。
程序
public class FindKNNInBST { public List<Integer> findKNN(TreeNode root, int target, int k) { List<Integer> knnList = new ArrayList<Integer>(); if (root == null || k <= 0) { return knnList; } Comparator<Integer> comp = new Comparator<Integer>() { @Override public int compare(Integer o1, Integer o2) { // compare abs value o1 = Math.abs(o1); o2 = Math.abs(o2); if (o1 < o2) { return 1; } else if (o1 > o2) { return -1; } return 0; } }; PriorityQueue<Integer> pq = new PriorityQueue<Integer>(k, comp); inorderTraversal(root, target, k, pq); for (Integer diff : pq) { knnList.add(target + diff); } return knnList; } private void inorderTraversal(TreeNode root, int target, int k, PriorityQueue<Integer> pq) { if (root == null) { return; } inorderTraversal(root.left, target, k, pq); if (pq.size() < k) { pq.add(root.val - target); } else { int diff = Math.abs(root.val - target); int maxDiff = Math.abs(pq.peek()); if (maxDiff > diff) { pq.poll(); pq.add(root.val - target); } } inorderTraversal(root.right, target, k, pq); } }