使用布尔矩阵从数据框中提取向量

Question

我有一个看起来像这样的数据框：

df = pd.DataFrame(np.random.rand(10,3), columns = ['col1', 'col2', 'col3'])

	col1	col2	col3
0	0.692154	0.286560	0.515904
1	0.798917	0.777593	0.971300

我还有另一个矩阵，它是一个布尔矩阵，看起来像这样：

b_matrix = pd.DataFrame(np.array([[0,1,1],
                       [1,1,0],
                       [0,0,1],
                       [0,1,0]]),                                             
                       columns = ['col1', 'col2', 'col3'],
                       index = ['input1', 'input2', 'input3', 'input4'])

	col1	col2	col3
input1	0	1	1
input2	1	1	0
input3	0	0	1
input4	0	1	0

所以这里的想法是我将提供一些输入，这将根据 b_matrix 进行检查，然后我将只返回 df 的相应列。例如如果输入是 input1 那么输出将是 df[['col2', 'col3']]:

	col2	col3
0	0.286560	0.515904
1	0.777593	0.971300

我可以想办法做到这一点，每次都要检查列名的静态列表，但我想知道是否有更直接的方法？

Answer 1

您可以使用 input to take user input and df.loc:

In [1076]: inp = input('User input:')
User input:input1

In [1077]: df[b_matrix.columns[b_matrix.loc[inp].eq(1)]]
Out[1077]: 
       col2      col3
0  0.179902  0.832655
1  0.444187  0.487146
2  0.879333  0.756792
3  0.870601  0.661337
4  0.082169  0.008669
5  0.190734  0.975966
6  0.839718  0.290976
7  0.862724  0.426222
8  0.581909  0.333300
9  0.949953  0.539106

如果选择input2:

In [1080]: inp = input('User input:')
User input:input2

In [1081]: df[b_matrix.columns[b_matrix.loc[inp].eq(1)]]
Out[1081]: 
       col1      col2
0  0.072600  0.179902
1  0.126708  0.444187
2  0.646533  0.879333
3  0.673643  0.870601
4  0.313205  0.082169
5  0.951917  0.190734
6  0.076799  0.839718
7  0.294087  0.862724
8  0.240569  0.581909
9  0.851999  0.949953

使用布尔矩阵从数据框中提取向量

Extracting vectors from a dataframe using a Boolean matrix

python

boolean

numpy

pandas