找到 K 个最接近的元素

题目要求

给定一个排序好的数组，两个整数 k 和 x，从数组中找到最靠近 x（两数之差最小）的 k 个数。返回的结果必须要是按升序排好的。如果有两个数与 x 的差值一样，优先选择数值较小的那个数

输入: [1, 5, 8, 9, 17, 20], k=4, x=10
输出: [5, 8, 9, 17]

思路分析

第一个大于等于目标值的元素位置

题目要求找到K个最接近目标值的元素，我们先不考虑K个，而是考虑1个的情况，也就是最接近目标值的那个元素。

如果数组里有某个元素恰好和目标值相等，那么问题就转变成了二分查找法，在数组中查找目标值。如果数组里没有目标值呢，如何找到距离目标值最近的元素呢？

不妨先找到第一个大于等于目标值的元素，找到了这个元素，也就容易找到距离目标值最近的那个元素了，分析到这里，我们先要实现一个函数，该函数可以从数组中找到第一个大于等于目标值的元素位置，下面是这个函数的示例代码

def find_eg_index(lst, start, end, target):
    """
    找到第一个大于等于target的元素的位置, 如果lst中最大的元素小于target
    则返回len(lst)
    :param lst:
    :param start:
    :param end:
    :param target:
    :return:
    """
    if start > end:
        return end + 1
    middle = (start + end)//2
    middle_value = lst[middle]
    if middle_value == target:
        return middle
    elif middle_value < target:
        return find_eg_index(lst, middle+1, end, target)
    else:
        return find_eg_index(lst, start, middle-1, target)

一定要记得考虑边界条件，如果目标值比数组最小值还小，那么函数返回0是应当的，但如果目标值比数组中最大值还大，该返回什么呢?为了概念的完整性，我这里返回len(lst), 这个索引是超出数组范围的，这表示无穷大，无穷大当然大于等于target

距离目标值最近的元素位置

前面的find_eg_index函数返回了第一个大于等于target的元素的位置eg_index，在此基础上，很容易就算出来距离target最近的元素的位置，所要考虑的逻辑分支如下

• eg_index == len(lst), target > max(lst), 数组最后一个元素最接近target
• eg_index == 0, 数组第一个元素最接近target
• target - lst[eg_index-1] <= lst[eg_index] - target, 说明lst[eg_index-1]更接近target, 反之lst[eg_index]更接近target

实现函数find_closest_index(lst, target) ，返回数组中距离target最近的元素的位置

def find_closest_index(lst, target):
    """
    寻找距离target最近的位置
    :param lst:
    :param target:
    :return:
    """
    eg_index = find_eg_index(lst, 0, len(lst)-1, target)
    if eg_index == len(lst):
        return eg_index - 1
    elif eg_index == 0:
        return 0
    else:
        if target - lst[eg_index-1] <= lst[eg_index] - target:
            return eg_index-1
        else:
            return eg_index

K个最接近的元素

在函数find_closest_index基础上继续思考，已经获得了closest_index， 只需要找到剩余的K-1个元素就可以了。

剩余的K-1个元素，一定分布在closest_index 的两侧，借鉴归并排序的思路，把closest_index 两侧的元素视为两个数组，分别从left_index(closest_index-1)和right_index(closest_index + 1)开始进行遍历比较，如果lst[left_index] <= lst[right_index]则lst[left_index] 就是要找的那K-1个元素里的一个，此时left_index -= 1，游标向左滑动，继续进行比较

这里要考虑边界情况，如果left_index 超出了数组的索引范围，该如何比较呢，一侧超出了索引范围，剩余的最近元素一定集中在另一侧，此时，可以直接从另一侧数组里直接取剩余数量的元素，但我认为这样程序就会写的很复杂，要判断两侧的索引是否超出范围，要写很多if语句。

不弱这样处理，如果索引超出了范围，就认为这个索引上的元素是无穷大，这样就不需要判断左右两个索引是否超出范围以及取另一侧的剩余最近元素了。

示例代码

def find_closet_ele(lst, target, k):
    closet_ele_lst = []
    closest_index = find_closest_index(lst, target)   # 距离target最近的元素的位置
    closet_ele_lst.append(lst[closest_index])

    ele_count = 1
    left_index = closest_index - 1
    right_index = closest_index + 1
    # 借鉴归并排序思路,向两侧滑动遍历
    while ele_count < k:
        left_ele = get_ele(lst, left_index)
        right_ele = get_ele(lst, right_index)
        # 哪边元素距离target更近,哪边就走一步
        if target - left_ele <= right_ele - target:
            closet_ele_lst.append(left_ele)
            left_index -= 1
        else:
            closet_ele_lst.append(right_ele)
            right_index += 1
        ele_count += 1

    closet_ele_lst.sort()
    return closet_ele_lst


def get_ele(lst, index):
    # 索引超出范围的,返回无穷大
    if index < 0 or index >= len(lst):
        return float("inf")
    return lst[index]

完整代码

def find_closet_ele(lst, target, k):
    closet_ele_lst = []
    closest_index = find_closest_index(lst, target)   # 距离target最近的元素的位置
    closet_ele_lst.append(lst[closest_index])

    ele_count = 1
    left_index = closest_index - 1
    right_index = closest_index + 1
    # 借鉴归并排序思路,向两侧滑动遍历
    while ele_count < k:
        left_ele = get_ele(lst, left_index)
        right_ele = get_ele(lst, right_index)
        # 哪边元素距离target更近,哪边就走一步,这里注意取绝对值
        if abs(target - left_ele) <= abs(right_ele - target):
            closet_ele_lst.append(left_ele)
            left_index -= 1
        else:
            closet_ele_lst.append(right_ele)
            right_index += 1
        ele_count += 1

    closet_ele_lst.sort()
    return closet_ele_lst


def get_ele(lst, index):
    # 索引超出范围的,返回无穷大
    if index < 0 or index >= len(lst):
        return float("inf")
    return lst[index]


def find_closest_index(lst, target):
    """
    寻找距离target最近的位置
    :param lst:
    :param target:
    :return:
    """
    eg_index = find_eg_index(lst, 0, len(lst)-1, target)
    if eg_index == len(lst):
        return eg_index - 1
    elif eg_index == 0:
        return 0
    else:
        if target - lst[eg_index-1] <= lst[eg_index] - target:
            return eg_index-1
        else:
            return eg_index


def find_eg_index(lst, start, end, target):
    """
    找到第一个大于等于target的元素的位置, 如果lst中最大的元素小于target
    则返回len(lst)
    :param lst:
    :param start:
    :param end:
    :param target:
    :return:
    """
    if start > end:
        return end + 1
    middle = (start + end)//2
    middle_value = lst[middle]
    if middle_value == target:
        return middle
    elif middle_value < target:
        return find_eg_index(lst, middle+1, end, target)
    else:
        return find_eg_index(lst, start, middle-1, target)


if __name__ == '__main__':
    lst = [1, 5, 8, 9, 17, 20]
    print(find_closet_ele(lst, 10, 4))