第 K 小的距离对

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

#
# 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
#
#
# @param nums int整型一维数组
# @param k int整型
# @return int整型
#
class Solution:
    """
    题目：
        https://www.nowcoder.com/practice/26849d7e31dd443c895a6921d67fa273?tpId=196&tqId=40558&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进行升序排序
        2. nums中任意数对的距离都在区间[0, max(nums) - min(nums)]中，所以第k小的距离对也处于该区间中，我们在该区间上进行二分查找。
        3. 对于当前数对距离mid，统计距离小于等于mid的数对总和cnt，如果cnt >= k，取right = mid - 1，否则left = mid + 1。当left > right时，退出循环，第k小的数对距离就是left。
        4. 对于cnt的计算：使用滑动窗口[i, j], i指向左边界，j指向右边界，我们右移i使得区间[i, j]所有数对满足小于等于mid，统计当前区间的数独j - i（j - i + 1个数组成j - i个数对），然后右移j，继续上述步骤进行统计，直至统计完所有区间。
    复杂度：
        时间复杂度：O(nlogn + n*logd)，排序算法：O(nlogn)，数对最大差值为d，在区间[0, d]上进行二分，每一次二分，都需要通过滑动窗口统计满足条件的数对，复杂度为：O(logd * n)，总体的复杂度为O(nlogn + n*logd)
        空间复杂度：O(logn)，排序算法的复杂度
    """

    def kthDistance(self, nums, k):
        # write code here
        nums.sort()

        left, right = 0, nums[-1] - nums[0]
        while left <= right:
            mid = left + (right - left) / 2

            cnt, i = 0, 0
            for j, num in enumerate(nums):
                while num - nums[i] > mid:
                    i += 1
                cnt += j - i

            if cnt >= k:
                right = mid - 1
            else:
                left = mid + 1

        return left


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

    # nums, k = [1, 3, 1, 5], 1

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

    nums, k = [6, 8, 5, 16, 3, 8, 2, 15, 10, 7, 1, 2], 20

    res = sol.kthDistance(nums, k)

    print res