【面試題】求兩個有序數組兩兩相加的值最小的K個數

題目：

有兩個大小都是k的數組A,B，它們元素的按非遞減有序排列，找出這樣的k個最小的(ai + bj) ,其中 0<= i,j < k，要求算法的時間複雜度和空間複雜度儘量低。

例如對於：

A = 1,2,3,4

B = 2,3,4,5

ai+bj的所有組合有4*4 = 16個，如下圖：

b\a 1   2   3   4

2   3   4   5   6

3   4   5   6   7

4   5   6   7   8

5   6   7   8   9

依次排序爲：3,4,4,5, 5,5,6,6， 6,6，7,7, 7,8,8,9 （共16個）

在舉一個例子：

A = 1,2,3,4

B = 20,30，40,50

ai+bj的所有組合有4*4 = 16個，如下圖：

b\a  1   2   3   4

20   21  22  23  24

30   31  32  33  34

40   41  42  43  44

50   51  52  53  54

依次排序爲：21,22,23,24,31,32,33,34,41,42,43,44,51,52,53,54（共16個）

考慮代碼實現，首先最小的必然是a0+b0，接下來是a1+b0、a0+b1中的小值，如果a1+b0小（第2個），接下來看a1+b1、a2+b0、a0+b1中哪個小……

一開始我用這種思路去思考，以爲可以找到O(K)時間複雜度的算法，最後發現不行。

方法一：

後來在網上看到了一份別人寫的用最小堆方法實現的代碼，可惜代碼有問題。其思路如下：

首先把a0+b0的結果放入堆中，此時堆中只有一個元素，自然滿足最小堆條件，然後開始出堆的操作，從堆裏面取出根節點（也就是最小的值），例如是a[i]+b[j]，則需要像最小堆中壓入a[i+1]b[j] 和 a[i]+b[j+1]，當然，要保證下標不越界，如果下標越界了則忽略，另外要保證已經壓入過堆中的組合（即使已經從堆中被取出了的）不再被壓入堆中。不段進行出堆、入堆的操作，重複K次，就得到了K個最小的組合值。

堆的最大深度爲logK,所以時間複雜度爲K*logK數量級。

方法二：

和同事導論得到另外一個解法，需要一個k長度的輔助數組記爲nIdxArray，一開始數組各項全部初始化爲0。然後用b[0]到b[k]分別加上a[nIdxArray[0]]到a[nIdxArray[k]]，遍歷一遍得到最小的值並得到對應的a數組下標i，更新nIdxArray[i]的值爲nIdxArray[i]+1。

然後繼續用b[0]到b[k]加上對應的a[nIdxArray[0]]到a[nIdxArray[k]]，遍歷得到其中的最小值，並更新nIdxArray對應項（+1操作）。

每得到一個數值需要比較K次，得到K個最小的值需要O(K*K)時間複雜度。

方法三：

暴力算法雖然是最慢的，但是是最容易寫出來，也是最可靠的。最重要的一點是，暴力算法的結果可以用來檢驗其他算法是否正確。

將ai+bj的所有值都保存起來，用一次快排搞定之後，前面K個就是最小的K個值。

三種方式的代碼實現：

// twoArrayKmin.cpp : Defines the entry point for the console application.
//

#define ARRAY_SIZE 10
#include <time.h>
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>

typedef struct _HEAP_ELEMENT
{
    int nIdxA;
    int nIdxB;
    int nSum;
    _HEAP_ELEMENT()
    {};
    _HEAP_ELEMENT(int i, int j, int n)
    {
        nIdxA = i;
        nIdxB = j;
        nSum  = n;
    }
}HEAP_ELEMENT;

int GetMin(int *nArrayA, int *nArrayB, int *nIdx, int nSize);
void Heap_insert(HEAP_ELEMENT element);
int GetHeapMin(HEAP_ELEMENT* pElement);
void checkThreeArrayEqueal(int *a, int *b, int *c, int nSize);
void Do();

int compare(const void* a, const void* b)
{
     return (*(int*)a - *(int*)b);
}

HEAP_ELEMENT g_heap[ARRAY_SIZE*ARRAY_SIZE+1];
HEAP_ELEMENT g_outArray[ARRAY_SIZE*ARRAY_SIZE+1];
int g_nHeapSize = 0; 
int g_nOutSize = 0; 

int main()
{
    srand(time(NULL));
    for (int i = 0; i < 1000; i++)
    {
        Do();
    }
    return 0;
}

void GetByQsort(int *nResultByQsort, int *nArrayA, int *nArrayB)
{
    for (int i = 0; i < ARRAY_SIZE; i++)
    {
        for (int j = 0; j < ARRAY_SIZE; j++)
        {
            nResultByQsort[i*ARRAY_SIZE + j] = nArrayA[i] + nArrayB[j];
        }
    }
    qsort(nResultByQsort, ARRAY_SIZE*ARRAY_SIZE, sizeof(int), compare);
    printf("get by qsort:\r\n");
    for (int i = 0; i < ARRAY_SIZE*ARRAY_SIZE; i++)
        printf("%02d  ", nResultByQsort[i]);
    printf("\r\n");
}

void GetByAssArray(int *nResultByMerge, int *nArrayA, int *nArrayB)
{
    int m, n;
    m = n = 0;
    int nIdx[ARRAY_SIZE] = {0};
    for (int i = 0; i < ARRAY_SIZE*ARRAY_SIZE; i++)
    {
        nResultByMerge[i] = GetMin(nArrayA, nArrayB, nIdx, ARRAY_SIZE);
    }
    printf("get by assistant array:\r\n");
    for (int i = 0; i < ARRAY_SIZE*ARRAY_SIZE; i++)
        printf("%02d  ", nResultByMerge[i]);
    printf("\r\n");
}

void GetByMinHeap(int *nResultByMinHeap, int *nArrayA, int *nArrayB)
{
    HEAP_ELEMENT element;
    g_nOutSize = 0;
    Heap_insert(HEAP_ELEMENT(0, 0, nArrayA[0] + nArrayB[0]));
    int num = 0;
    for (int i = 0; i < ARRAY_SIZE*ARRAY_SIZE; i++)
    {
        nResultByMinHeap[i] = GetHeapMin(&element);
        //printf("No. %d: %d %d %d \r\n", ++num, element.nIdxA, element.nIdxB, element.nSum);
        if (element.nIdxA+1 < ARRAY_SIZE)
            Heap_insert(HEAP_ELEMENT(element.nIdxA+1, element.nIdxB, 
            nArrayA[element.nIdxA+1]+nArrayB[element.nIdxB]));
        if (element.nIdxB+1 < ARRAY_SIZE)
            Heap_insert(HEAP_ELEMENT(element.nIdxA, element.nIdxB+1, 
            nArrayA[element.nIdxA]+nArrayB[element.nIdxB+1]));
        assert(g_nHeapSize < ARRAY_SIZE*ARRAY_SIZE-i);
    }
    assert(g_nHeapSize == 0);
    printf("get by min heap:\r\n");
    for (int i = 0; i < ARRAY_SIZE*ARRAY_SIZE; i++)
        printf("%02d  ", nResultByMinHeap[i]);
    printf("\r\n");
}

void Do()
{
    //array A: 03  09  15  15  19
    //array B: 03  14  18  19  21
    //array A: 01  03  06  08  08  12  12  14  14  27
    //array B: 05  08  10  10  10  11  13  18  24  28
    int nArrayA[ARRAY_SIZE];// = {3,  9,  15,  15,  19};
    int nArrayB[ARRAY_SIZE];// = {3,  14,  18,  19,  21};

    int nResultByQsort[ARRAY_SIZE*ARRAY_SIZE];
    int nResultByMerge[ARRAY_SIZE*ARRAY_SIZE];
    int nResultByMinHeap[ARRAY_SIZE*ARRAY_SIZE];

    for (int i = 0; i < ARRAY_SIZE; i++)
    {
        nArrayA[i] = rand()%30;
        nArrayB[i] = rand()%30;
    }

    qsort(nArrayA, ARRAY_SIZE, sizeof(int), compare);
    qsort(nArrayB, ARRAY_SIZE, sizeof(int), compare);

    printf("array A: ");
    for (int i = 0; i < ARRAY_SIZE; i++)
        printf("%02d  ", nArrayA[i]);
    printf("\r\n");

    printf("array B: ");
    for (int i = 0; i < ARRAY_SIZE; i++)
        printf("%02d  ", nArrayB[i]);
    printf("\r\n");

    // get by qsort
    GetByQsort(nResultByQsort, nArrayA, nArrayB);

    // get by assistant array
    GetByAssArray(nResultByMerge, nArrayA, nArrayB);

    // get by min heap
    GetByMinHeap(nResultByMinHeap, nArrayA, nArrayB);

    checkThreeArrayEqueal(nResultByMerge, nResultByMinHeap, nResultByQsort, ARRAY_SIZE*ARRAY_SIZE);

}

void checkThreeArrayEqueal(int *a, int *b, int *c, int nSize)
{
    for (int i = 0; i < nSize; i++)
    {
        assert(a[i] == b[i]);
        assert(b[i] == c[i]);
    }
}

int GetMin(int *nArrayA, int *nArrayB, int *nIdx, int nSize)
{
    int iFind = 0;
    int nMin = nArrayA[nSize-1] + nArrayB[nSize-1] + 1;
    for (int i = 0; i < nSize; i++)
    {
        if (nIdx[i] < nSize && nArrayA[i] + nArrayB[nIdx[i]] < nMin)
        {
            nMin = nArrayA[i] + nArrayB[nIdx[i]];
            iFind = i;
        }
    }
    nIdx[iFind]++;
    return nMin;
}

void SwapElement(int i, int j)
{
    HEAP_ELEMENT element = g_heap[i];
    g_heap[i] = g_heap[j];
    g_heap[j] = element;
}

void Heap_insert(HEAP_ELEMENT element)
{
    //check the element not in the heap
    for (int i = 1; i <= g_nHeapSize; i++)
    {
        if (g_heap[i].nIdxA == element.nIdxA &&
            g_heap[i].nIdxB == element.nIdxB &&
            g_heap[i].nSum == element.nSum)
            return;
    }

    //check the element not in output array
    for (int i = 0; i < g_nOutSize; i++)
    {
        if (g_outArray[i].nIdxA == element.nIdxA &&
            g_outArray[i].nIdxB == element.nIdxB &&
            g_outArray[i].nSum == element.nSum)
            return;
    }

    int n = ++g_nHeapSize;
    g_heap[n] = element;
    while (n/2)
    {
        if (g_heap[n].nSum < g_heap[n/2].nSum)
        {
            SwapElement(n, n/2);
            n = n/2;
        }
        else
            break;
    }
}

void Heapify(int i)
{
    int left = 2*i;
    int right = 2*i+1;
    int nSmall;

    if (g_nHeapSize < left)
        return;

    nSmall = left;
    if ( (right <= g_nHeapSize) && (g_heap[right].nSum < g_heap[left].nSum) )
        nSmall = right;

    if (g_heap[i].nSum >= g_heap[nSmall].nSum)
        SwapElement(i, nSmall);

    Heapify(nSmall);
}

int GetHeapMin(HEAP_ELEMENT* pElement)
{
    *pElement = g_heap[1];
    g_outArray[g_nOutSize++] = g_heap[1];

    g_heap[1] = g_heap[g_nHeapSize--];
    Heapify(1);

    return pElement->nSum;
}