csv文件的快速排序？

Question

我对Python还不是很熟悉，主要是想学它来做数据分析。我有一个CSV文件，里面的内容大概是这样的（ID，类别）：

67512367,0
67567,5
89789789,5
...
...
876289347,10
638, 10
...
...
98723489,20
3828909, 20
...
...<going upto>
78789789, 200
978789, 200

我想做的是按照类别来排序这些ID。最后我希望得到的结果看起来像这样：

list_5 = [67567, 89789789, .., ]
list_10 = [876289347, 638, ...]
list_200 = [78789789, 978789, ...]

问题是这个CSV文件大约有150万条记录。所以我现在是用izip和csv读取器来加载它们，像这样：

data = izip(csv.reader(open("data.csv", "rb")))

然后我把这些记录放在data里，可以简单地用以下方式来遍历：

for i in data:
    print i
    #print i[0][0] # for ids
    #print i[0][1] # for category

现在，我知道我可以用if-else/elif结构来检查i[0][1]==5，然后把i[0][0]（ID）添加到一个列表里，但这样做似乎会很慢，而且我的列表很大。

我在想有没有其他更优雅的方法（也许可以用itertools？）来根据第二列的值（类别）来对ID进行“分桶”。

Answer 1

调用排序两次。

sortedCats = sorted(data, key=attrgetter('category_name'))
sortedIds = sorted(sortedCats, key=attrgetter('id'))

这样做是有效的，因为Python使用的排序算法叫做timsort，第二次排序会利用第一次排序时形成的分组，从而减少工作量。

Answer 2

你可以使用 itertools.groupby 这个工具：

# test_big.csv was 1.74GB
def test():
    from itertools import groupby
    bucketized_grouped_keys = {}
    with open('test_big','r') as f:
        for key, group in groupby(f,lambda T: T.split(',')[1].rstrip('\n')):
                bucketized_grouped_keys.update({key:group})
    print(bucketized_grouped_keys.keys())

if __name__ == '__main__':
    import cProfile
    cProfile.run('test()','test.profile')
    import pstats
    stats = pstats.Stats('test.profile')
    stats.strip_dirs().sort_stats('time').print_stats()

输出结果是：

['47', '44', '2', '42', '49']
Thu Aug 07 10:55:39 2014    test.profile

         445620949 function calls in 239.002 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
148540313   95.738    0.000  178.467    0.000 csv_test.py:6(<lambda>)
        1   60.535   60.535  239.002  239.002 csv_test.py:2(test)
148540313   55.128    0.000   55.128    0.000 {method 'split' of 'str' objects}
148540313   27.601    0.000   27.601    0.000 {method 'rstrip' of 'str' objects}

        1    0.000    0.000    0.000    0.000 {open}
        1    0.000    0.000  239.002  239.002 <string>:1(<module>)
        5    0.000    0.000    0.000    0.000 {method 'update' of 'dict' objects
}
        1    0.000    0.000    0.000    0.000 {method 'keys' of 'dict' objects}
        1    0.000    0.000    0.000    0.000 {method 'disable' of '_lsprof.Prof
iler' objects}

关键函数是：

lambda T: T.split(',')[1].rstrip('\n')

这和下面这个是一样的：

def T(item):
    return item.split(',')[1].rstrip('\n')

要访问这些分组：

category_name = '42'
bucketized_grouped_keys[category_name]

Answer 3

既然你说你“主要是想学习Python来做数据分析”，那么你一定要看看pandas这个工具，这样你就能有更好的工具来玩了。（当然，了解如何从零开始构建类似pandas的工具也是有用的。但根据我的经验，即使使用pandas，你也有足够的机会来锻炼你的Python技能，而且弄清楚如何做一些实际的事情要比重新实现基本功能更有趣。）

你可以使用read_csv把文件读入一个数据框（就像Excel表格一样）：

>>> import pandas as pd
>>> df = pd.read_csv("group.csv", names=["ID", "category"])
>>> df
          ID  category
0   67512367         0
1      67567         5
2   89789789         5
3  876289347        10
4        638        10
5   98723489        20
6    3828909        20
7   78789789       200
8     978789       200

然后使用groupby来建立一个类别到ID的字典：

>>> {k: v.tolist() for k,v in df.groupby("category")["ID"]}
{0: [67512367], 200: [78789789, 978789], 10: [876289347, 638], 20: [98723489, 3828909], 5: [67567, 89789789]}

虽然你也可以直接对groupby对象进行很多操作（比如计算统计数据等），所以老实说，我并不常需要一个实际的索引列表。更常见的情况是我只是想“对每个组执行这个操作”，但具体情况可能因人而异。