如何在Pandas DataFrame中按特定规则提取唯一行？

原始数据 #

我们有如下Pandas DataFrame：

import pandas as pd
df = pd.DataFrame(
    {
        "ID": [1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4],
        "color": ["red","blue","green","yellow","red","blue","green","yellow","red","blue","green","yellow","red","blue","green","yellow"]
    }
)

每个ID对应red、blue、green、yellow四种颜色，需要提取的目标结果为：ID1对应red、ID2对应blue、ID3对应green、ID4对应yellow。

无效尝试 #

使用df.drop_duplicates()按单一列去重无法满足需求：

• 按ID去重：df.drop_duplicates("ID")，结果所有ID对应的颜色均为red；
• 按color去重：df.drop_duplicates("color")，结果所有颜色对应的ID均为1；
  keep参数仅能反转保留的首个/尾个值，无法实现自定义匹配规则。

可行解决方案 #

方案一：分组循环提取 #

这是初始想到的较繁琐实现方式：

series = []
for i, (_, sub) in enumerate(df.groupby("ID")):
    series.append(sub.iloc[i])

out = pd.concat(series, axis=1).transpose().reset_index(drop=True)

方案二：匹配元组筛选 #

更简洁高效的实现，适合ID和color列唯一值长度相等的场景：

zipped = zip(df["ID"].unique(), df["color"].unique())
mask = df[["ID", "color"]].agg(tuple, 1).isin(zipped)

df[mask]

注：zip会以最短序列为准，但在本场景中ID和color的唯一值长度始终相等，因此无需顾虑。

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