剑指 Offer II 061. 和最小的 k 个数对-大顶堆法

剑指 Offer II 061. 和最小的 k 个数对

给定两个以升序排列的整数数组 nums1 和 nums2 , 以及一个整数 k 。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。

请找到和最小的 k 个数对 (u1,v1), (u2,v2) … (uk,vk) 。

示例 1:

输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [1,2],[1,4],[1,6]
解释: 返回序列中的前 3 对数：
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

示例 2:

输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数：
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

示例 3:

输入: nums1 = [1,2], nums2 = [3], k = 3
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]

解题代码如下：

/**
 * Return an array of arrays of size *returnSize.
 * The sizes of the arrays are returned as *returnColumnSizes array.
 * Note: Both returned array and *columnSizes array must be malloced, assume caller calls free().
 */
 /* 大顶堆 注意点：数据类型需要严格一致，long就都用long*/
typedef struct MyHeap {
    int cap;
    long **data;
    int index;
}MyHeap;
bool HeapCmp(long a, long b) {
    if (a > b) {
        return true;
    }
    return false;
}
void HeapSwap(long *src1, long *src2, long *dst1, long *dst2) {
    long tmp1 = *src1;
    long tmp2 = *src2;
    *src1 = *dst1;
    *src2 = *dst2;

    *dst1 = tmp1;
    *dst2 = tmp2;
}
void HeapUp(MyHeap *heap) {
    int child = heap->index - 1;
    int parent = (child - 1) / 2;

    while (child > 0) {
        long sum1 =  heap->data[child][0] + heap->data[child][1];
        long sum2 =  heap->data[parent][0] + heap->data[parent][1];
        if (HeapCmp(sum1, sum2)) {
            HeapSwap(&heap->data[child][0],  &heap->data[child][1],  &heap->data[parent][0],  &heap->data[parent][1]);
        } else {
            break;
        }
        child = parent;
        parent = (child - 1) / 2;
    }
}

void HeapDown(MyHeap *heap) {
    int parent = 0;
    int child = parent * 2 + 1;

    while (child < heap->cap) {
        if (child + 1 < heap->cap && HeapCmp(heap->data[child + 1][0] + heap->data[child + 1][1], heap->data[child][0] + heap->data[child][1])) {
            child = child + 1;
        }
        long sum1 =  heap->data[child][0] + heap->data[child][1];
        long sum2 =  heap->data[parent][0] + heap->data[parent][1];
        if (HeapCmp(sum1, sum2)) {
            HeapSwap(&heap->data[child][0],  &heap->data[child][1],  &heap->data[parent][0],  &heap->data[parent][1]);
        } else {
            break;
        }
        parent = child;
        child = parent * 2 + 1;
    }
}

void add(MyHeap *heap, int x1, int x2) {
    if (heap->index < heap->cap) {
        
        heap->data[heap->index][0] = x1;
        heap->data[heap->index][1] = x2;
        heap->index++;
        HeapUp(heap);      
    } else {
        long sum1 =  (long)x1 + (long)x2;
        long sum2 =  heap->data[0][0] + heap->data[0][1];
        if (sum1 < sum2) {
            long x = (long)x1;
            long y = (long)x2;
            HeapSwap(&x, &y, &heap->data[0][0], &heap->data[0][1]);
            HeapDown(heap);
        }
    }
}
int** kSmallestPairs(int* nums1, int nums1Size, int* nums2, int nums2Size, int k, int* returnSize, int** returnColumnSizes){
    MyHeap *heap = (MyHeap *)malloc(sizeof(MyHeap));
    heap->cap = k;
    heap->index = 0;
    heap->data = (long **)malloc(sizeof(long *) * k);
    for (int i = 0; i < k; i++) {
        heap->data[i] = (long *)malloc(sizeof(long) * 2);
    }
   // printf("111\n");
    for (int i = 0; i < nums1Size; i++) {
        if (heap->index == heap->cap) {
            long sum1 =  (long)nums1[i] + (long)nums2[0];
            long sum2 =  heap->data[0][0] + heap->data[0][1];
            if (sum1 > sum2) {
                break; /* 剪枝，如果最小的都大，就没必要继续了，直接退出 */
            }
        }
        for (int j = 0; j < nums2Size; j++) {
            add(heap, nums1[i], nums2[j]);
        }
    }
    *returnSize = heap->index;
    *returnColumnSizes = (int *)malloc(sizeof(int) * heap->index);
    for (int i = 0; i < heap->index; i++) {
        (*returnColumnSizes)[i] = 2;
    }

    int **rcs = (int *)malloc(sizeof(int *) * heap->index);
    for (int i = 0; i < heap->index; i++) {
        rcs[i] = (int *)malloc(sizeof(int) * 2);
        rcs[i][0] = heap->data[i][0];
        rcs[i][1] = heap->data[i][1];
    }
    return rcs;
}