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);
	}
}