在Python中对数据进行子集提取

Question

我想在我写的Python代码中使用类似于R语言的subset命令。

这是我的数据：

col1    col2    col3    col4    col5
100002  2006    1.1 0.01    6352
100002  2006    1.2 0.84    304518
100002  2006    2   1.52    148219
100002  2007    1.1 0.01    6292
10002   2006    1.1 0.01    5968
10002   2006    1.2 0.25    104318
10002   2007    1.1 0.01    6800
10002   2007    4   2.03    25446
10002   2008    1.1 0.01    6408

我想根据col1和col2的内容来筛选数据。（在col1中，唯一的值是100002和10002，而在col2中，唯一的值是2006、2007和2008。）

在R中可以使用subset命令来做到这一点，Python中有没有类似的东西呢？

Answer 1

因为我对R语言不太熟悉，也不太了解你提到的这个子集命令是怎么工作的，所以我建议你看看Python中的itertools库里的groupby功能。这个功能可以根据你给定的一个函数的输出值来分组。如果你想了解更多，可以查看groupby的文档。

groups = []
uniquekeys = []
data = sorted(data, key=keyfunc)
for k, g in groupby(data, keyfunc):
    groups.append(list(g))      # Store group iterator as a list
    uniquekeys.append(k)

这样你就得到了你的子集。不过要注意，返回的值并不是完整的列表，而是迭代器。

我假设你的值是按行返回的。

Answer 2

在R语言中，subset()的功能和Python中的filter()很相似。正如参考资料所提到的，这个功能会在列表推导式中被隐式使用，所以写代码时最简洁明了的方式可能是

[ item for item in items if item.col2 == 2006 ] 

比如说，如果你的数据行放在一个叫items的可迭代对象里。

Answer 3

虽然基于迭代器的答案也没问题，但如果你在使用numpy数组（正如你提到的那样），其实有更好更快的方法来选择数据：

import numpy as np
data = np.array([
        [100002, 2006, 1.1, 0.01, 6352],
        [100002, 2006, 1.2, 0.84, 304518],
        [100002, 2006, 2,   1.52, 148219],
        [100002, 2007, 1.1, 0.01, 6292],
        [10002,  2006, 1.1, 0.01, 5968],
        [10002,  2006, 1.2, 0.25, 104318],
        [10002,  2007, 1.1, 0.01, 6800],
        [10002,  2007, 4,   2.03, 25446],
        [10002,  2008, 1.1, 0.01, 6408]    ])

subset1 = data[data[:,0] == 100002]
subset2 = data[data[:,0] == 10002]

这样可以得到

子集1：

array([[  1.00002e+05,   2.006e+03,   1.10e+00, 1.00e-02,   6.352e+03],
       [  1.00002e+05,   2.006e+03,   1.20e+00, 8.40e-01,   3.04518e+05],
       [  1.00002e+05,   2.006e+03,   2.00e+00, 1.52e+00,   1.48219e+05],
       [  1.00002e+05,   2.007e+03,   1.10e+00, 1.00e-02,   6.292e+03]])

子集2：

array([[  1.0002e+04,   2.006e+03,   1.10e+00, 1.00e-02,   5.968e+03],
       [  1.0002e+04,   2.006e+03,   1.20e+00, 2.50e-01,   1.04318e+05],
       [  1.0002e+04,   2.007e+03,   1.10e+00, 1.00e-02,   6.800e+03],
       [  1.0002e+04,   2.007e+03,   4.00e+00, 2.03e+00,   2.5446e+04],
       [  1.0002e+04,   2.008e+03,   1.10e+00, 1.00e-02,   6.408e+03]])

如果你事先不知道第一列的唯一值，可以使用numpy.unique1d或者内置的set函数来找出来。

补充：我刚意识到你想选择的是两列的唯一组合数据……在这种情况下，你可以这样做：

col1 = data[:,0]
col2 = data[:,1]

subsets = {}
for val1, val2 in itertools.product(np.unique(col1), np.unique(col2)):
    subset = data[(col1 == val1) & (col2 == val2)]
    if np.any(subset):
        subsets[(val1, val2)] = subset

（我把这些子集存储在一个字典里，字典的键是组合的元组……当然还有其他（更好，具体要看你做什么）方法可以做到这一点！）