1 #!/usr/bin/env python 2 # coding=utf-8 3 4 ''' 5 蓄水池抽样法。从不知道总数的样本中随机抽取k个,要求只遍历一次样本。 6 ''' 7 8 import sys 9 import random 10 11 if __name__ == '__main__': 12 '''命令行参数:输入文件,抽取多少行,输出文件。输出文件不指定时,默认为输出到控制台''' 13 if len(sys.argv) < 3: 14 print "input [infile] [lines] [outfile=console]" 15 infile = sys.argv[1] 16 k = int(sys.argv[2]) 17 if len(sys.argv) >= 4: 18 f_handler = open(sys.argv[3], 'w') 19 sys.stdout = f_handler 20 rect = [] 21 lineNo = 0 22 with open(infile, 'r') as f_in: 23 for i in range(k): 24 line = f_in.readline() 25 if line: 26 lineNo += 1 27 rect.append(line.strip('\n')) 28 else: 29 break 30 while True: 31 line = f_in.readline() 32 if not line: 33 break 34 lineNo += 1 35 rnd = random.randint(1, lineNo + 1) 36 if rnd < k: 37 rect[rnd] = line.strip('\n') 38 for ele in rect: 39 print ele