Leetcode日练5 #658 Find K closest element

#658 Find K Closest Elements

Given a sorted integer array arr, two integers k and x, return the k closest integers to x in the array. The result should also be sorted in ascending order.

An integer a is closer to x than an integer b if:

• |a - x| < |b - x|, or
• |a - x| == |b - x| and a < b

解题思路（pseudocode）：

1.用常规的二分法，找到arr里与x相等或离x最近的整数。

2.以这个数为轴，其 idx ± k 后框出备选范围[idx-k: idx+k+1]；edge case，可能idx-k后小于0，或者idx+k+1后超出len(arr)。这两种情况取0或者len(arr)-1。取备选范围的所有数与x的差值，以及其原来的idx信息，新生成一个list of tuples。

3.再把这个list根据差值从小到大sorted，组成一个dict，key是差值，val是list of对应的idx。

4.循环这个dict，直到取够k个idx为止。key是从小到大排列的，所以当目前val内的idx个数小于k时，则全部取完。若大于k，则只取部分。取idx的规则是看val这个list有没有被x断开。如果x在val之外或者是头尾，则从x近侧取值。比如，在左侧则从小到大取，在右侧就从大到小取。如果被断开了，就是从断裂处的左侧从大往小取，取完了断裂处左侧的idx再取右侧的。

global arr, k, x, left, right, diff_list

arr = [1,3,3,4,5,7,7,8,8,8]
k = 6
x = 6


# find the closest integer: mid is its index
def find_the_closest_integer(arr, x):
    left, right = 0, len(arr) - 1
    while right >= left:
        mid = (left + right) // 2
        # print(mid)
        if arr[mid] == x:
            break
        elif arr[mid] > x:  # x is on the left
            right = mid - 1
        else:  # x is on the right
            left = mid + 1
    print(mid)
    # either mid is the index of the first item after the gap if x is not in the list
    # or the index of the first x in the list
    return mid


# narrow down the scope to find k closest integers
# diff_list is a list of tuple with index and diff
def potential_scope(arr, k, x, mid):
    diff_list = []

    if mid - k >= 0:
        start = mid - k
    else:
        start = 0
    if mid + k + 1 <= len(arr) - 1:
        end = mid + k + 1
    else:
        end = len(arr) - 1
    # print(start, end)
    for i in range(start, end+1):
        diff_list.append((i, abs(arr[i] - x)))
    # print(diff_list)
    return diff_list


# rearrange the scope in a dict based on diff starting from zero
def pool(diff_list):
    pool = sorted(diff_list, key=lambda i: i[1])
    # print(pool)
    pre_diff = []
    pool_dict = {}
    for idx, diff in pool:
        pre_diff.append(diff)

    for i in sorted(set(pre_diff)):  # for each diff
        # print(i)
        pool_dict[i] = []
        for idx, diff in pool:
            if i == diff:
                pool_dict[i].append(idx)

    print(pool_dict)
    return pool_dict


# pick integers from the pool dict
def extract(pool_dict, mid, k):
    num = 0     # counter of qualified indices
    selected = []   # store qualified indices
    for key