最接近的K个元素

# -*- coding: utf-8 -*-

import bisect


#
# 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
#
#
# @param nums int整型一维数组
# @param k int整型
# @param x int整型
# @return int整型一维数组
#
class Solution:
    """
    题目：
        https://www.nowcoder.com/practice/b4d7edc45759453e9bc8ab71f0888e0f?tpId=196&tqId=40554&rp=1&ru=/exam/oj&qru=/exam/oj&sourceUrl=%2Fexam%2Foj%3FjudgeStatus%3D3%26page%3D1%26pageSize%3D50%26search%3D%26tab%3D%25E7%25AE%2597%25E6%25B3%2595%25E7%25AF%2587%26topicId%3D196&difficulty=undefined&judgeStatus=3&tags=&title=
    算法：
        1. 在升序数组nums中，二分查找大于等于x的下标right，与小于等于x的下标left
        2. 以区间[left, right]为中心，向外扩展，记录满足条件的元素：
            |a-x| < |b-x|
            |a-x| = |b-x| 时 a < b
    复杂度：
        时间复杂度：O(logn + k)
        空间复杂度：O(1)
    """

    def closestElement(self, nums, k, x):
        # write code here
        left = bisect.bisect_left(nums, x)

        if left > 0 and abs(nums[left - 1] - x) <= abs(nums[left] - x): # bisect.bisect_left返回的是nums中大于等于x的最小下标，所以这里最接近x的有可能是nums[left], 也可能是nums[left - 1]
            i, j = left - 1, left
        else:
            i, j = left, left + 1

        n, res = len(nums), []
        while i >= 0 and j < n and k > 0:
            if abs(nums[i] - x) <= abs(nums[j] - x):
                res.insert(0, nums[i])
                i -= 1
            else:
                res.append(nums[j])
                j += 1
            k -= 1

        if i >= 0 and k > 0:
            res = nums[i - k + 1:i + 1] + res

        if j < n and k > 0:
            res += nums[j: j + k]

        return res


if __name__ == "__main__":
    sol = Solution()

    nums, k, x = [1, 2, 3, 4, 5], 4, 4

    # nums, k, x = [1, 2, 3, 4, 5], 4, 3

    # nums, k, x = [1, 2, 3, 4], 4, 0

    # nums, k, x = [1, 3, 8, 9], 1, 5

    res = sol.closestElement(nums, k, x)

    print res