第 K 小的距离对

import java.util.*;


public class Solution {
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     *
     * @param nums int整型ArrayList
     * @param k int整型
     * @return int整型
     */
    public int kthDistance (ArrayList<Integer> nums, int k) {
        // write code here
        Collections.sort(nums);

        int n = nums.size();
        int left = 0;
        int right = nums.get(n - 1) - nums.get(0);
        int mid;
        // 二分
        while (left <= right) {
            mid = left + (right - left) / 2;
            // left最终指向第一个cnt大于等于k的数
            // if(check(nums, k, mid)){
            if (check(nums, k, mid)) {
                left = mid + 1;
            } else {
                right = mid - 1;
            }
        }
        return left;
    }

    private boolean check(ArrayList<Integer> nums, int k, int target) {
        int n = nums.size();
        // 小于等于target的个数
        int cnt = 0;
        // 双指针
        for (int i = 0; i < n; i++) {
            // 再次二分 进行优化
            int left = i + 1, right = n - 1;
            int mid;
            while (left <= right) {
                mid = left + ((right - left) >> 1);
                // left最终指向第一个满足条件(nums.get(left)-nums.get(i) > target)的位置(从左往右)
                if (nums.get(mid) - nums.get(i) <= target) {
                    left = mid + 1;
                } else {
                    right = mid - 1;
                }
            }
            cnt += left - i - 1;
            if (cnt >= k) {
                return false;
            }
        }

        return true;
    }
}

# -*- 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