python 基于小顶堆实现随机抽样

起因:之前用蓄水池抽样,算法精简,但直观性很差。 所以这次采用了简单的,为没一个行,赋值一个随机值,然后取 最大的K个作为,随机样本。

基本思路:为每一个行(record,记录,实体) 赋一个random值。 每个map取一个Top K 值。 由于是求topk,可以设置一个reduce,再求 Top K

map阶段:
一般思路, 读取所有的,sort,取前 10000.------>运行时发现巨慢。 之后换成了 小顶堆 求TopK。速度就很快了。


import heapq
import random

class Url(object):

    def __init__(self, url, c, lkt, r):
        self.url = url
        self.c = c
        self.lkt = lkt
        self.r =r

    def __cmp__(self, x):
        return cmp(self.r, x.r)

    def __str__(self):
        return ("%s " * 4).strip() % (self.url, self.c, self.lkt, self.r)


class TopKHeap(object):
    def __init__(self, k):
        self.k = k
        self.data = []

    def Push(self, elem):
        if len(self.data) < self.k:
            heapq.heappush(self.data, elem)
        else:
            topk_small = self.data[0]
            if elem > topk_small:
                heapq.heapreplace(self.data, elem)

    def TopK(self):
        return [x for x in reversed([heapq.heappop(self.data) for x in xrange(len(self.data))])]



import sys
import topk

K = 100000
tkh = topk.TopKHeap(K)
#test = []

#url,count,link_found_time,random
for line in sys.stdin:
    line = line.strip()
    arr = line.split('x01')
    if len(arr) != 4:
        continue

    arr[3] = float(arr[3])
    url = topk.Url(arr[0], arr[1], arr[2], arr[3])
    tkh.Push(url)

    #test.append(url)

for e in tkh.TopK():
    print e

#test code, test success
#print '---------------------------------'
#for e in sorted(test, reverse=True)[0: 10]: (这种类型的sorded太慢了)
#    print e