时间戳

首页 > 解决方案 > 带有列表理解的 Pandas DataFrame 条件选择

python - 带有列表理解的 Pandas DataFrame 条件选择

问题描述

我有一个包含 15 列的数据框，名为 0,1,2,...,14。我想编写一个方法来接收这些数据和一个长度为 15 的向量。我希望它返回基于我传递的这个向量有条件地选择的数据帧。例如，传递的数据是 data_，传递的向量是 v_ 我想生成：

data[(data[0] == v_[0]) & (data[1] == v_[1]) & ... & (data[14] == v_[14])]

但是我希望该方法灵活，例如我可以传入名为 0、...、99 和长度为 99 的向量的 100 列的数据帧。我的问题是我不知道如何巧妙地以[(data[0] == v_[0]) & (data[1] == v_[1]) & ... & (data[14] == v_[14])]编程方式创建来解释“＆“ 符号。同样，如果有人给我一种方法，可以将多个在“and”或“or”上填充了 True 和 False 的 NxM 矩阵合并到单个 MxN 矩阵中，我也会很满意。

非常感谢！

标签： pythonpandasnumpy

解决方案

你可以试试这个：

def custom_filter(data, v):
    if len(data.columns) == len(v):
        # If data has the same number of columns
        #  as v has elements
        mask = (data == v).all(axis=1)
    else:
        # If they have a different length, we'll need to subset
        #  the data first, then create our mask
        #  This attempts to susbet the dataframe by assuming columns
        #    0 .. len(v) - 1 exist as columns, and will throw an error
        #    otherwise
        colnames = list(range(len(v)))
        mask = (data[colnames] == v).all(axis=1)
    
    return data.loc[mask, :]

df = pd.DataFrame({
    "F": list("hiadsfin"),
    0: list("aaaabbbb"),
    1: list("cccdddee"),
    2: list("ffgghhij")
})
v = ["a", "c", "f"]

df
   F  0  1  2  H
0  h  a  c  f  1
1  i  a  c  f  2
2  a  a  c  g  3
3  d  a  d  g  4
4  s  b  d  h  5
5  f  b  d  h  6
6  i  b  e  i  7
7  n  b  e  j  8

custom_filter(df, v)
   F  0  1  2  H
0  h  a  c  f  1
1  i  a  c  f  2

请注意，使用此函数，如果列数与向量的长度完全匹配v，则无需确保列标记为0, 1, 2, ..., len(v)-1。但是，如果您的列多于 v 的元素，则需要确保这些列的子集标记为0, 1, 2, ..., len(v)-1. If v` 比数据框中的列长，这将引发错误。

推荐阅读