最近在hadoop上写了一个分布式的KNN,其实本地的KNN也是需要用到的一个东西——从N个距离中选出K个最小的,从而投票得出预测的类别。
这个问题看似很简单,但要做好,还是有难度的。
方案1
n个数字排序,然后截取前k个作为结果。
时间复杂度:O(nlog n)
空间复杂度:O(n)
时间复杂度还好,但是空间复杂度是关于输入规模线性的就不好了,因为当需要预测的数据很多,多到单机内存放不下时,就得考虑外排序了,这样设计起来,整个系统变复杂很多了,不是一个很好的设计。
方案2
使用插入排序维护一个大小为K+1的数组,比如K=3,然后依次插入5,3,1,4,6,0,7,那么数组的变化情况如下:
下面的[x]表示“未赋值”:
(1) 5 [x] [x] [x]
(2) 3 5 [x] [x]
(3) 1 3 5 [x]
(4) 1 3 4 [5](目前最大的5溢出了,我们只关注前k=3个数字)
(5) 1 3 4 [6](很明显,新插入的6也处于“溢出位”)
(6) 0 1 3 [4](4被挤到了“溢出位”)
(6) 0 1 3 [7](注意“溢出位”是被后续的数据覆盖掉的)
所以最终的输出为前k=3个最小的数字:0, 1, 3
我把它封装成一个C++的模板类,这么简单的算法,相信理解的人很容易将其转换成别的语言。
复杂度分析:
1. 时间复杂度:O(KN),当N > 2^K时,很明显这个方案是很有优势的,而一般K很小,这个条件比较容易满足。
2. 空间复杂度:O(K),注意K在KNN算法里是固定大小的,不随输入数据的规模改变
最后,注意我为什么说数组的大小是K+1而不是K,这是因为编程实现的时候比较方便,不用特判去处理边界。
// TopKMinElem.h
#ifndef __TOK_K_MIN_ELEM__H
#define __TOK_K_MIN_ELEM__H
#include <vector>
using namespace std;
template<typename KeyType, typename ElemType>
class TopKMinElem
{
private:
int index, capacity;
vector<pair<KeyType, ElemType> > data;
public:
TopKMinElem(int k)
: index(0), capacity(k), data(k+1) {}
void insert(KeyType k, ElemType e) {
int pos = index - 1;
while (pos >= 0 && k < data[pos].first) {
data[pos+1] = data[pos];
--pos;
}
if (pos+1 < capacity)
data[pos+1] = make_pair(k, e);
index = min(index+1, capacity);
}