如何根据条件从NumPy Matrix获取行的子集?

时间:2016-03-24 17:22:34

标签: python numpy matrix

如何返回与给定条件匹配的NumPy Matrix的一组行?

这是一个Numpy Matrix对象

>>> X

matrix([['sunny', 'hot', 'high', 'FALSE'],
        ['sunny', 'hot', 'high', 'TRUE'],
        ['overcast', 'hot', 'high', 'FALSE'],
        ['rainy', 'mild', 'high', 'FALSE'],
        ['rainy', 'cool', 'normal', 'FALSE'],
        ['rainy', 'cool', 'normal', 'TRUE'],
        ['overcast', 'cool', 'normal', 'TRUE'],
        ['sunny', 'mild', 'high', 'FALSE'],
        ['sunny', 'cool', 'normal', 'FALSE'],
        ['rainy', 'mild', 'normal', 'FALSE'],
        ['sunny', 'mild', 'normal', 'TRUE'],
        ['overcast', 'mild', 'high', 'TRUE'],
        ['overcast', 'hot', 'normal', 'FALSE'],
        ['rainy', 'mild', 'high', 'TRUE']], 
       dtype='|S8')

我想获取第一列值为'rainy'的所有行的集合,所以它尝试了这个

>>> X[X[:,0]=='rainy']

matrix([['rainy', 'rainy', 'rainy', 'rainy', 'rainy']], 
       dtype='|S8')

但我想要一个这样的输出

matrix([['rainy', 'mild', 'high', 'FALSE'],
        ['rainy', 'cool', 'normal', 'FALSE'],
        ['rainy', 'cool', 'normal', 'TRUE'],
        ['rainy', 'mild', 'normal', 'FALSE'],
        ['rainy', 'mild', 'high', 'TRUE']], 
       dtype='|S8')

应该怎么做?

2 个答案:

答案 0 :(得分:3)

>>> X[(X[:, 0] == 'rainy').ravel(), :]
matrix([['rainy', 'mild', 'high', 'FALSE'],
        ['rainy', 'cool', 'normal', 'FALSE'],
        ['rainy', 'cool', 'normal', 'TRUE'],
        ['rainy', 'mild', 'normal', 'FALSE'],
        ['rainy', 'mild', 'high', 'TRUE']], 
       dtype='|S8')

如果你看一下比较的结果:

>>> X[:, 0] == 'rainy'
array([[False],
       [False],
       [False],
       [ True],
       [ True],
       [ True],
       [False],
       [False],
       [False],
       [ True],
       [False],
       [False],
       [False],
       [ True]], dtype=bool)

需要使用ravel

将其展平为矢量
(X[:, 0] == 'rainy').ravel()
array([False, False, False,  True,  True,  True, False, False, False,
        True, False, False, False,  True], dtype=bool)

对于其他约束,这有效:

X[(X[:, 0] == 'rainy').ravel() & (X[:, 1] == 'cool').ravel(), :]
matrix([['rainy', 'cool', 'normal', 'FALSE'],
        ['rainy', 'cool', 'normal', 'TRUE']], 
       dtype='|S8')

答案 1 :(得分:1)

这样做的方法不止一种。

foo = np.where(X[:, 0] == 'rainy') # get the index
X[foo, :]                          # The result you want.