根据存储在字典中的条件从Pandas数据框中选择数据。

Question

我有一个Pandas数据框，其中包含大量的变量。可以简化为：

tempDF = pd.DataFrame({ 'var1': [12,12,12,12,45,45,45,51,51,51],
                        'var2': ['a','a','b','b','b','b','b','c','c','d'],
                        'var3': ['e','f','f','f','f','g','g','g','g','g'],
                        'var4': [1,2,3,3,4,5,6,6,6,7]})

如果我想选择数据框的一个子集（例如，var2='b'和var4=3），我会使用以下代码：

tempDF.loc[(tempDF['var2']=='b') & (tempDF['var4']==3),:]

然而，如果匹配的条件存储在字典中，例如：是否可能选择数据框的子集？

tempDict = {'var2': 'b','var4': 3}

变量名不应预定义，并且字典中包含的变量数量是可变的，这一点非常重要。

我已经思考了一段时间，因此非常感谢任何建议。

- user1718097

Answer 1

您可以评估一系列条件，它们不仅仅是相等性。

df = tempDF
d = tempDict

# `repr` returns the string representation of an object.    
>>> df[eval(" & ".join(["(df['{0}'] == {1})".format(col, repr(cond)) 
       for col, cond in d.iteritems()]))]
   var1 var2 var3  var4
2    12    b    f     3
3    12    b    f     3

看一下这里的eval作用：

conditions = " & ".join(["(df['{0}'] == {1})".format(col, repr(cond)) 
       for col, cond in d.iteritems()])

>>> conditions
"(df['var4'] == 3) & (df['var2'] == 'b')"

>>> eval(conditions)
0    False
1    False
2     True
3     True
4    False
5    False
6    False
7    False
8    False
9    False
dtype: bool

这里是另一个使用等式约束的例子：

>>> eval(" & ".join(["(df['{0}'] == {1})".format(col, repr(cond)) 
                      for col, cond in d.iteritems()]))
d = {'var2': ('==', "'b'"),
     'var4': ('>', 3)}

>>> df[eval(" & ".join(["(df['{0}'] {1} {2})".format(col, cond[0], cond[1]) 
       for col, cond in d.iteritems()]))]
   var1 var2 var3  var4
4    45    b    f     4
5    45    b    g     5
6    45    b    g     6

另一种选择是使用query方法：

qry = " & ".join('{0} {1} {2}'.format(k, cond[0], cond[1]) for k, cond in d.iteritems())

>>> qry
"var4 > 3 & var2 == 'b'"

>>> df.query(qry)
   var1 var2 var3  var4
4    45    b    f     4
5    45    b    g     5
6    45    b    g     6

- Alexander

---

非常感谢您的回复，Alexander。不幸的是，我无法使您提供的第一个解决方案工作（AttributeError: 'dict' object has no attribute 'iteritems'）。 - user1718097

你使用的是哪个版本的Python？对于Python 3，应该使用d.items()而不是d.iteritems()。这个问题上没有Python 3标签。 - Alexander

你说得对，我应该包含一个Python 3标签。 - user1718097

这是三个答案中发表的其中一个，时间相距仅几分钟。所有答案都非常有效，并且按照要求进行操作。此外，我从未想到过任何解决方案 - 因此在这个过程中我学到了很多东西。然而，我选择了Anton Protopopov的答案（使用列表推导和掩码），因为它更符合我的思路。 - user1718097

Answer 2

你可以使用列表推导式为每个条件创建掩码，然后将它们连接起来，先转换成数据框再使用 all 方法：

In [23]: pd.DataFrame([tempDF[key] == val for key, val in tempDict.items()]).T.all(axis=1)
Out[23]:
0    False
1    False
2     True
3     True
4    False
5    False
6    False
7    False
8    False
9    False
dtype: bool

然后您可以使用该掩码来切片数据框：

mask = pd.DataFrame([tempDF[key] == val for key, val in tempDict.items()]).T.all(axis=1)

In [25]: tempDF[mask]
Out[25]:
   var1 var2 var3  var4
2    12    b    f     3
3    12    b    f     3

- Anton Protopopov

---

谢谢Anton。解决方案很好用。我需要理解列表推导和掩码的使用。 - user1718097

这是三个答案中的一个，它们在几分钟内发布。所有答案都非常有效，并且完全按照要求执行。此外，我从未想过任何解决方案-因此在这个过程中我学到了很多。然而，我选择了这个作为选定的答案，因为它似乎更符合我的思维方式。 - user1718097

Answer 3

这里有一种从tempDict构建条件的方法。

In [25]: tempDF.loc[pd.np.all([tempDF[k] == tempDict[k] for k in tempDict], axis=0), :]
Out[25]:
   var1 var2 var3  var4
2    12    b    f     3
3    12    b    f     3

或者使用query来获取更易读的类似查询的字符串。

In [33]: tempDF.query(' & '.join(['{0}=={1}'.format(k, repr(v)) for k, v in tempDict.iteritems()]))
Out[33]:
   var1 var2 var3  var4
2    12    b    f     3
3    12    b    f     3

In [34]: ' & '.join(['{0}=={1}'.format(k, repr(v)) for k, v in tempDict.iteritems()])
Out[34]: "var4==3 & var2=='b'"

- Zero

---

谢谢 John。解决方案非常好。但是此刻我想不出为什么需要 pd.np.all()，我需要再想一下…… - user1718097

这是三个答案中的一个，它们在几分钟内发布。所有答案都非常有效，并且完全按照要求执行。此外，我自己从未想过任何解决方案-因此在这个过程中我学到了很多。然而，我选择了Anton Protopopov的答案（使用列表推导和掩码），因为它更符合我的思维方式。 - user1718097

Answer 4

这是我个人工具中的一个函数，它接受单个值或列表进行子集操作：

def subsetdict(df, sdict):
    subsetter_list = [df[i].isin([j]) if not isinstance(j, list) else df[i].isin(j) for i, j in sdict.items()]
    subsetter = pd.concat(subsetter_list, axis=1).all(1)
    return df.loc[subsetter, :]

- ghart