658. 找到 K 个最接近的元素
题目描述
输入一个已按升序排列的数组 arr,以及整数 k 和 x。
比较两个数 a、b 哪个更靠近 x:先比 |a - x| 与 |b - x|,较小的更近;二者相等时,数值较小的那个算更近。
按这个规则从 arr 中选出最靠近 x 的 k 个元素,并以升序返回。
输入:arr = [1,2,3,4,5], k = 4, x = 3 输出:[1,2,3,4]
解释:离 3 最近的是 3、2、4,再在 1 和 5 之间按平局规则选较小的 1。
输入:arr = [1,1,2,3,4,5], k = 4, x = -1 输出:[1,1,2,3]
解释:x 比所有元素都小,取最前面的 4 个。
数据范围
1 <= k <= arr.length1 <= arr.length <= 10^4arr按升序排列。-10^4 <= arr[i], x <= 10^4
解题思路:二分查找窗口左端点
答案一定是一段长度为 k 的连续窗口 arr[lo:lo+k],只需二分窗口左端点 lo,范围是 [0, n-k]。比较窗口左端 arr[mid] 和窗口外右边第一个数 arr[mid+k]:如果 x - arr[mid] > arr[mid+k] - x,说明右边那个数更近,窗口必须右移,lo = mid + 1;否则左端不差于右边,最优左端点在 mid 或更左。
为什么正确
最优窗口的左端点始终在 [lo, hi] 中。
复杂度分析
时间 O(log(n-k) + k),空间 O(1)(不计输出)。
易错点
比较时不能用绝对值 |x-arr[mid]| 与 |arr[mid+k]-x|:当 x 在窗口两端同一侧时绝对值会给出错误方向。用带符号的差即可,相等时偏向左边。
参考代码
Python、Java、C++、Go、JavaScript、TypeScript 六种语言,逻辑一致。
class Solution:
def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
lo, hi = 0, len(arr) - k
while lo < hi:
mid = (lo + hi) // 2
if x - arr[mid] > arr[mid + k] - x:
lo = mid + 1
else:
hi = mid
return arr[lo:lo + k]class Solution {
public List<Integer> findClosestElements(int[] arr, int k, int x) {
int lo = 0, hi = arr.length - k;
while (lo < hi) {
int mid = (lo + hi) / 2;
if (x - arr[mid] > arr[mid + k] - x) {
lo = mid + 1;
} else {
hi = mid;
}
}
List<Integer> res = new ArrayList<>();
for (int i = lo; i < lo + k; i++) res.add(arr[i]);
return res;
}
}class Solution {
public:
vector<int> findClosestElements(vector<int>& arr, int k, int x) {
int lo = 0, hi = (int)arr.size() - k;
while (lo < hi) {
int mid = (lo + hi) / 2;
if (x - arr[mid] > arr[mid + k] - x) {
lo = mid + 1;
} else {
hi = mid;
}
}
return vector<int>(arr.begin() + lo, arr.begin() + lo + k);
}
};func findClosestElements(arr []int, k int, x int) []int {
lo, hi := 0, len(arr)-k
for lo < hi {
mid := (lo + hi) / 2
if x-arr[mid] > arr[mid+k]-x {
lo = mid + 1
} else {
hi = mid
}
}
return arr[lo : lo+k]
}var findClosestElements = function(arr, k, x) {
let lo = 0, hi = arr.length - k;
while (lo < hi) {
const mid = Math.floor((lo + hi) / 2);
if (x - arr[mid] > arr[mid + k] - x) {
lo = mid + 1;
} else {
hi = mid;
}
}
return arr.slice(lo, lo + k);
};function findClosestElements(arr: number[], k: number, x: number): number[] {
let lo = 0, hi = arr.length - k;
while (lo < hi) {
const mid = Math.floor((lo + hi) / 2);
if (x - arr[mid] > arr[mid + k] - x) {
lo = mid + 1;
} else {
hi = mid;
}
}
return arr.slice(lo, lo + k);
}其他解法
按距离排序没有利用数组有序;双指针利用“答案是连续区间”从两端删元素;二分查找直接定位长为 k 的窗口的左端点,最快。
解法 2:双指针收缩窗口 时间 O(n),空间 O(1)
结果一定是一个长为 k 的连续区间。从整个数组 [lo, hi] 开始,每次删掉两端中离 x 更远的一个。距离相等时删右端,因为题目要求优先保留较小的数。区间长度缩到 k 时返回。
思路直观,没有二分的边界细节。代价是要删 n − k 次,k 很小时比二分慢。
class Solution:
def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
lo, hi = 0, len(arr) - 1
while hi - lo + 1 > k:
if x - arr[lo] <= arr[hi] - x:
hi -= 1
else:
lo += 1
return arr[lo:hi + 1]class Solution {
public List<Integer> findClosestElements(int[] arr, int k, int x) {
int lo = 0, hi = arr.length - 1;
while (hi - lo + 1 > k) {
if (x - arr[lo] <= arr[hi] - x) hi--;
else lo++;
}
List<Integer> ans = new ArrayList<>();
for (int i = lo; i <= hi; i++) ans.add(arr[i]);
return ans;
}
}class Solution {
public:
vector<int> findClosestElements(vector<int>& arr, int k, int x) {
int lo = 0, hi = arr.size() - 1;
while (hi - lo + 1 > k) {
if (x - arr[lo] <= arr[hi] - x) hi--;
else lo++;
}
return vector<int>(arr.begin() + lo, arr.begin() + hi + 1);
}
};func findClosestElements(arr []int, k int, x int) []int {
lo, hi := 0, len(arr)-1
for hi-lo+1 > k {
if x-arr[lo] <= arr[hi]-x {
hi--
} else {
lo++
}
}
return arr[lo : hi+1]
}var findClosestElements = function(arr, k, x) {
let lo = 0, hi = arr.length - 1;
while (hi - lo + 1 > k) {
if (x - arr[lo] <= arr[hi] - x) hi--;
else lo++;
}
return arr.slice(lo, hi + 1);
};function findClosestElements(arr: number[], k: number, x: number): number[] {
let lo = 0, hi = arr.length - 1;
while (hi - lo + 1 > k) {
if (x - arr[lo] <= arr[hi] - x) hi--;
else lo++;
}
return arr.slice(lo, hi + 1);
}解法 3:按距离排序 时间 O(n log n),空间 O(n)
把元素按与 x 的距离从小到大排序,距离相同则较小的数在前。取前 k 个,再按数值升序排列后返回。
直接把题意翻译成排序规则,不依赖输入有序。代价是 O(n log n) 的排序和额外空间。
class Solution:
def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
byDist = sorted(arr, key=lambda v: (abs(v - x), v))
return sorted(byDist[:k])class Solution {
public List<Integer> findClosestElements(int[] arr, int k, int x) {
List<Integer> byDist = new ArrayList<>();
for (int v : arr) byDist.add(v);
byDist.sort((a, b) -> {
int da = Math.abs(a - x), db = Math.abs(b - x);
return da != db ? Integer.compare(da, db) : Integer.compare(a, b);
});
List<Integer> ans = new ArrayList<>(byDist.subList(0, k));
Collections.sort(ans);
return ans;
}
}class Solution {
public:
vector<int> findClosestElements(vector<int>& arr, int k, int x) {
vector<int> byDist(arr);
sort(byDist.begin(), byDist.end(), [x](int a, int b) {
int da = abs(a - x), db = abs(b - x);
return da != db ? da < db : a < b;
});
vector<int> ans(byDist.begin(), byDist.begin() + k);
sort(ans.begin(), ans.end());
return ans;
}
};func findClosestElements(arr []int, k int, x int) []int {
abs := func(v int) int {
if v < 0 {
return -v
}
return v
}
byDist := append([]int(nil), arr...)
sort.Slice(byDist, func(i, j int) bool {
di, dj := abs(byDist[i]-x), abs(byDist[j]-x)
if di != dj {
return di < dj
}
return byDist[i] < byDist[j]
})
ans := byDist[:k]
sort.Ints(ans)
return ans
}var findClosestElements = function(arr, k, x) {
const byDist = [...arr].sort((a, b) => Math.abs(a - x) - Math.abs(b - x) || a - b);
return byDist.slice(0, k).sort((a, b) => a - b);
};function findClosestElements(arr: number[], k: number, x: number): number[] {
const byDist = [...arr].sort((a, b) => Math.abs(a - x) - Math.abs(b - x) || a - b);
return byDist.slice(0, k).sort((a, b) => a - b);
}