Python实现分层抽样

首先说明一下我的需求。result_33.txt的文件中有一些号码标签及分类的标记。具体如下

中国农业银行 银行
招商银行信用卡 信用卡
门窗 无关

我想做的是从每一个类标记中随机抽出1000个标签。如果该类标签下的样本数不足1000,则全部抽取。然后将抽取的结果保存到另一个文件中。

具体代码如下(尽可能减少内存占用量,但是文件扫描次数太多):

import random
if __name__ == '__main__':
    data={}
    with open("result_33.txt",'r',encoding='utf-8-sig') as r_dict,\
        open('check_result_33.txt','w',encoding='utf-8') as w_dict:
        for line in r_dict:
            items=line.strip('\n').split()
            data[items[1]]=data.get(items[1],0)+1
        for key,value in data.items():
            r_dict.seek(0)  # 返回文件开始位置
            if value<1000:
                for line in r_dict:
                    if key==line.strip('\n').split()[1]:
                        w_dict.write(line)
            else:
                rand_list=list(range(value))
                random.shuffle(rand_list)
                rand_list=iter(rand_list)
                for line in r_dict:
                    if key==line.strip('\n').split()[1] and next(rand_list)<1000:
                        w_dict.write(line)

 

你可能感兴趣的:(Python,分层抽样)