如何基于0/1矩阵筛选含指定ID字符串的Pandas DataFrame行

我来帮你搞定这个问题！你遇到的场景是需要从带有各种前缀后缀的ID列中提取纯净ID，再对照日期对应的0/1矩阵来筛选DataFrame行，用循环遍历确实效率低还容易出错，咱们用Pandas的向量化操作来高效解决。

先回顾一下你的示例数据 #

首先先把你给出的示例代码放出来，方便后续操作：

import pandas as pd

# 示例DataFrame df1
df1 = pd.DataFrame({ 
    'Date':['2021-01-01', '2021-01-01', '2021-01-01', '2021-01-02', '2021-01-02', '2021-01-03', '2021-01-03', '2021-01-04', '2021-01-05', '2021-01-05', '2021-01-05', '2021-01-05'], 
    'ID':['R: ID1', 'ID2_P', 'L_ID2', 'ID1.I', 'RZ:ID3', 'ID1', 'P ID2', 'ID3 4K', 'ID1', 'U_ID1', 'ID2:PD', 'ID3=2D'], 
    'Ratings':[9.0, 8.0, 5.0, 3.0, 2, 3, 6, 5, 4, 10, 3, 6]
})

# 0/1矩阵m1
m1 = pd.DataFrame({ 
    'Date':['2021-01-01', '2021-01-02', '2021-01-03', '2021-01-04', '2021-01-05'], 
    'ID1':[0,0,1,1,1], 
    'ID2':[1,1,1,1,0], 
    'ID3':[0,0,0,0,1]
})

步骤1：从df1的ID列提取纯净的ID标识 #

你的ID列里有各种额外字符（比如R: 、_P、.I等），我们需要先提取出核心的ID1/ID2/ID3这类标识。用正则表达式配合str.extract可以轻松实现：

# 提取ID列中的纯ID（匹配ID+数字的模式）
df1['Cleaned_ID'] = df1['ID'].str.extract(r'(ID\d+)')

执行后，df1会新增一列Cleaned_ID，每一行对应纯净的ID值，比如R: ID1对应ID1，ID2_P对应ID2。

步骤2：转换m1为长格式并筛选有效条件 #

当前m1是宽格式（ID作为列名），我们需要把它转成长格式，这样每一行对应一个Date-Cleaned_ID的组合，再筛选出值为1的有效组合：

# 将m1从宽格式转为长格式
m1_long = m1.melt(id_vars='Date', var_name='Cleaned_ID', value_name='Flag')
# 只保留Flag为1的有效组合
valid_pairs = m1_long[m1_long['Flag'] == 1]

valid_pairs里就存储了所有需要保留的（日期，纯净ID）组合，比如2021-01-01和ID2、2021-01-03和ID1等。

步骤3：合并数据得到筛选后的df2 #

最后我们把df1和valid_pairs做内连接，只保留同时匹配日期和纯净ID的行，再清理掉临时列即可：

# 内连接df1和有效组合，筛选符合条件的行
df2 = df1.merge(valid_pairs, on=['Date', 'Cleaned_ID'], how='inner')
# 删除临时生成的Cleaned_ID和Flag列
df2 = df2.drop(columns=['Cleaned_ID', 'Flag'])
# 重置索引（可选，让索引更规整）
df2 = df2.reset_index(drop=True)

验证结果 #

现在输出df2，就和你想要的目标DataFrame完全一致了：

Date       ID  Ratings
0  2021-01-01   ID2_P      8.0
1  2021-01-01   L_ID2      5.0
2  2021-01-03      ID1      3.0
3  2021-01-03    P ID2      6.0
4  2021-01-05      ID1      4.0
5  2021-01-05   U_ID1     10.0
6  2021-01-05  ID3=2D      6.0

这种方法用Pandas的原生向量化操作代替循环，不仅代码更简洁，处理大型DataFrame时的效率也会高很多哦！

内容的提问来源于stack exchange，提问作者fjurt