【数组】找到 K 个最接近的元素

题目描述

给定一个 排序好 的数组arr ，两个整数 k 和 x ，从数组中找到最靠近 x（两数之差最小）的 k 个数。返回的结果必须要是按升序排好的。

整数 a 比整数 b 更接近 x 需要满足：

• |a - x| < |b - x| 或者
• |a - x| == |b - x| 且 a < b

示例 1：

输入：arr = [1,2,3,4,5], k = 4, x = 3
输出：[1,2,3,4]

解题思路 

1. 找到绝对值最小的元素的下标 newIndex，直接遍历数组；
2. 使用newIndex作为起点，从左向右、从右向左添加元素，优先添加绝对值小的元素。

代码实现如下：

import java.util.LinkedList;
import java.util.List;

class Solution1 {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {

        int left = 0;
        int right = 0;

        int newIndex = findClosestIndex(arr, x);


        List<Integer> res = new LinkedList<>();


        res.add(arr[newIndex]);

        left = newIndex - 1;
        right = newIndex + 1;
        while (res.size() < k && left >= 0 && right < arr.length) {
            int absLeft = Math.abs(arr[left] - x);
            int absRight = Math.abs(arr[right] - x);
            if (absLeft <= absRight) {
                res.add(0, arr[left]);
                left--;
            } else {
                res.add(arr[right]);
                right++;
            }
        }

        while (res.size() < k && left >= 0) {
            res.add(0, arr[left]);
            left--;
        }

        while (res.size() < k && right < arr.length) {
            res.add(arr[right]);
            right++;
        }
        return res;
    }

    private int findClosestIndex(int[] arr, int x) {

        int minIndex = 0;
        int minAbs = Math.abs(x - arr[0]);

        for (int i = 1; i < arr.length; i++) {
            if (minAbs > Math.abs(x - arr[i])) {
                minAbs = Math.abs(x - arr[i]);
                minIndex = i;
            }
        }
        return minIndex;
    }

    public static void main(String[] args) {
        Solution1 solution = new Solution1();
        System.out.println(solution.findClosestElements(new int[]{1, 1, 1, 10, 10, 10}, 1, 9));
    }
}

---

这里可以进行优化，使用二分查找解决，为了能实现二分查找，代码改造后如下：

import java.util.*;

class Solution {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {

        int left = 0;
        int right = 0;

        int newIndex = findClosestIndex(arr, x);

        if (newIndex - 1 >= 0) {
            if (Math.abs(arr[newIndex] - x) >= Math.abs(arr[newIndex - 1] - x)) {
                newIndex--;
            }
        }

        if (newIndex + 1 < arr.length) {
            if (Math.abs(arr[newIndex] - x) > Math.abs(arr[newIndex + 1] - x)) {
                newIndex++;
            }
        }


        List<Integer> res = new LinkedList<>();


        res.add(arr[newIndex]);

        left = newIndex - 1;
        right = newIndex + 1;
        while (res.size() < k && left >= 0 && right < arr.length) {
            int absLeft = Math.abs(arr[left] - x);
            int absRight = Math.abs(arr[right] - x);
            if (absLeft <= absRight) {
                res.add(0, arr[left]);
                left--;
            } else {
                res.add(arr[right]);
                right++;
            }
        }

        while (res.size() < k && left >= 0) {
            res.add(0, arr[left]);
            left--;
        }

        while (res.size() < k && right < arr.length) {
            res.add(arr[right]);
            right++;
        }
        return res;
    }

    private int findClosestIndex(int[] arr, int x) {

        int left = 0;
        int right = arr.length - 1;

        while (left < right) {

            int mid = (right - left) / 2 + left;
            if (arr[mid] < x) {
                left = mid + 1;
            } else if (arr[mid] > x) {
                right = mid - 1;
            } else {
                right = mid - 1;
            }
        }
        return left;
    }


    public static void main(String[] args) {
        Solution solution = new Solution();
        System.out.println(solution.findClosestElements(new int[]{1, 1, 1, 10, 10, 10}, 1, 9));
        System.out.println(solution.findClosestElements(new int[]{1, 3}, 1, 2));
        System.out.println(solution.findClosestElements(new int[]{1, 5,10}, 1, 4));
    }
}

---

还可以通过另外一种实现方式，直接排序处理：

• 对每个元素的绝对值做升序排序；
• 从0到k-1加入到结果列表中；
• 对结果进行排序；
• 返回结果。

这块代码很简洁，代码如下：

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

class Solution2 {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {
        List<Integer> list = new ArrayList<>(arr.length);
        for (int v : arr) {
            list.add(v);
        }

        Collections.sort(list, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                int v = Math.abs(o1 - x) - Math.abs(o2 - x);
                return v == 0 ? o1 - o2 : v;
            }
        });

        List<Integer> res = new ArrayList<>(list.subList(0, k));
        Collections.sort(res);
        return res;
    }


    public static void main(String[] args) {
        Solution2 solution = new Solution2();
        System.out.println(solution.findClosestElements(new int[]{1, 1, 1, 10, 10, 10}, 1, 9));
    }
}

总结

这块代码还可以进一步优化耗时，如果有更加高效、简洁的实现，欢迎回复。