C语言 | Leetcode C语言题解之第373题查找和最小的K对数字

题目：

题解：

#define MIN(a, b) ((a) > (b) ? (b) : (a))

int** kSmallestPairs(int* nums1, int nums1Size, int* nums2, int nums2Size, int k, int* returnSize, int** returnColumnSizes) {
    if (nums1Size == 0 || nums2Size == 0 || k <= 0) {
        *returnSize = 0;
        return NULL;
    }

    /*二分查找第 k 小的数对和的大小*/
    int left = nums1[0] + nums2[0];
    int right = nums1[nums1Size - 1] + nums2[nums2Size - 1];
    int pairSum = right;
    while (left <= right) {
        int mid = left + ((right - left) >> 1);
        long long cnt = 0;
        int start = 0;
        int end = nums2Size - 1;
        while (start < nums1Size && end >= 0) {
            if (nums1[start] + nums2[end] > mid) {
                end--;
            } else {
                cnt += end + 1;
                start++;
            }
        }
        if (cnt < k) {
            left = mid + 1;
        } else {
            pairSum = mid;
            right = mid - 1;
        }
    }

    int ** ans = (int **)malloc(sizeof(int *) * k);
    *returnColumnSizes = (int *)malloc(sizeof(int) * k);
    int currSize = 0;
    int pos = nums2Size - 1;
    /*找到小于目标值 pairSum 的数对*/
    for (int i = 0; i < nums1Size; i++) {
        while (pos >= 0 && nums1[i] + nums2[pos] >= pairSum) {
            pos--;
        }
        for (int j = 0; j <= pos && k > 0; j++, k--) {
            ans[currSize] = (int *)malloc(sizeof(int) * 2);
            ans[currSize][0] = nums1[i];
            ans[currSize][1] = nums2[j];
            (*returnColumnSizes)[currSize] = 2;
            currSize++;
        }
    }

    /*找到等于目标值 pairSum 的数对*/
    pos = nums2Size - 1;
    for (int i = 0; i < nums1Size && k > 0; i++) {
        int start1 = i;
        while (i < nums1Size - 1 && nums1[i] == nums1[i + 1]) {
            i++;
        }
        while (pos >= 0 && nums1[i] + nums2[pos] > pairSum) {
            pos--;
        }
        int start2 = pos;
        while (pos > 0 && nums2[pos] == nums2[pos - 1]) {
            pos--;
        }
        if (nums1[i] + nums2[pos] != pairSum) {
            continue;
        }
        int count = (int) MIN(k, (long) (i - start1 + 1) * (start2 - pos + 1));
        for (int j = 0; j < count && k > 0; j++, k--) {
            ans[currSize] = (int *)malloc(sizeof(int) * 2);
            ans[currSize][0] = nums1[i];
            ans[currSize][1] = nums2[pos];
            (*returnColumnSizes)[currSize] = 2;
            currSize++;
        }
    }
    *returnSize = currSize;
    return ans;
}