【面试题】求两个有序数组两两相加的值最小的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
= element;
while (n/2)
{
if (g_heap
.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;
}