python - 如果列中的总和等于 n，则将 n 行作为输入并返回列名的函数

Question

我有一个大DataFrame的结构如下：

import pandas as pd

df = pd.DataFrame({'name1': [1, 0, 1,1],
                   'name2': [0, 0, 0,1],
                   'name3': [1, 1, 1,1],
                   'namen': [0, 0, 0,0]},
                  index=['label1', 'label2', 'label3', 'labeln'])
>>> df
      name1 name2 name3 name4
label1  1     0     1      1
label2  0     0     0      1
label3  1     1     1      1
label4  0     0     0      0

我正在尝试构建一个函数，该函数将n行名称作为参数求和所有列中的值，如果这些列的总和等于n ，则返回我的列名称。

例如，使用 label1、label2 和 label3 作为输入，我想获得以下输出：

def common_terms(*nargs):
   the function...

>>> common_terms(label1, label2, label3)
(name4)

或者

>>> common_terms(label1, label3)
(name1, name3)

我对在 Python 中构建函数知之甚少，但我的头脑确实停留在这一点上。你能帮我进步吗？

Answer 1

过滤行loc并测试每列是否全部，1然后过滤index：Series

def common_terms(*nargs):
   i = df.loc[list(nargs)].all()
   return i.index[i].tolist()

print (common_terms('label1', 'label2', 'label3'))
['namen']

print (common_terms('label1','label3'))
['name1', 'namen']