如何基于两列值以最少分组数高效分组DataFrame并映射回原表

假设我们有如下Pandas DataFrame：

import pandas as pd
test = [["Red",10],["Red",5],["Blue",10],["Green",10],["Blue",10],["Red",15]]
df = pd.DataFrame(test,columns=['Color','Price'])

输出结果：

Color  Price
0    Red     10
1    Red      5
2   Blue     10
3  Green     10
4   Blue     10
5    Red     15

需求是：基于'Color'或'Price'列构建分组，找到能生成最少分组数的方式，最终得到如下汇总结果：

Group    Price                     Color
0      0  [5, 15]                     [Red]
1      1     [10]  [Red, Blue, Green, Blue]

同时需要将Group标签映射回原数据表，得到：

Color  Price  Group
0    Red     10      1
1    Red      5      0
2   Blue     10      1
3  Green     10      1
4   Blue     10      1
5    Red     15      0

你尝试过分别按Color和Price单独分组：

• 按Color分组结果：

df_color_group = df.groupby(['Color'],as_index=False)['Price'].agg(list).reset_index()
df_color_group.rename({"index":"Group"},axis=1,inplace=True)

输出：

Group  Color        Price
0      0   Blue     [10, 10]
1      1  Green         [10]
2      2    Red  [10, 5, 15]

• 按Price分组结果：

df_price_group = df.groupby(['Price'],as_index=False)['Color'].agg(list).reset_index()
df_price_group.rename({"index":"Group"},axis=1,inplace=True)

输出：

Group  Price                     Color
0      0      5                     [Red]
1      1     10  [Red, Blue, Green, Blue]
2      2     15                     [Red]

但不知道如何得到预期的跨维度合并分组结果。

你的需求本质是找到数据中的连通分量：将每一行视为一个节点，如果两行的Color相同或Price相同，则节点之间建立连接，最终每个连通分量就是一个分组，这样能得到最少的分组数。

方法一：用NetworkX快速实现（简洁高效） #

import pandas as pd
import networkx as nx

# 初始化原始数据
test = [["Red",10],["Red",5],["Blue",10],["Green",10],["Blue",10],["Red",15]]
df = pd.DataFrame(test,columns=['Color','Price'])

# 创建图对象
G = nx.Graph()
# 添加所有行作为节点（用行索引作为节点ID）
G.add_nodes_from(df.index)

# 为同一Color的行添加连接边
for color in df['Color'].unique():
    indices = df[df['Color'] == color].index.tolist()
    # 同一Color的所有节点两两连边
    for i in range(len(indices)):
        for j in range(i+1, len(indices)):
            G.add_edge(indices[i], indices[j])

# 为同一Price的行添加连接边
for price in df['Price'].unique():
    indices = df[df['Price'] == price].index.tolist()
    for i in range(len(indices)):
        for j in range(i+1, len(indices)):
            G.add_edge(indices[i], indices[j])

# 获取所有连通分量，为每个节点分配Group编号
groups = list(nx.connected_components(G))
group_mapping = {}
for group_id, nodes in enumerate(groups):
    for node in nodes:
        group_mapping[node] = group_id

# 将Group标签映射回原DataFrame
df['Group'] = df.index.map(group_mapping)

# 生成目标汇总表
summary_df = df.groupby('Group').agg(
    Price=('Price', list),
    Color=('Color', list)
).reset_index()

方法二：纯Pandas实现（无需额外库） #

如果不想依赖第三方库，可以通过多次合并相同维度的组来实现：

import pandas as pd

test = [["Red",10],["Red",5],["Blue",10],["Green",10],["Blue",10],["Red",15]]
df = pd.DataFrame(test,columns=['Color','Price'])
df['Group'] = df.index  # 初始每个行单独为一组

# 循环合并相同Color的组
while True:
    # 获取每个Color对应的最小Group编号
    color_min_group = df.groupby('Color')['Group'].min().reset_index()
    df = df.merge(color_min_group, on='Color', suffixes=('', '_new'))
    # 没有变化则退出循环
    if df['Group'].equals(df['Group_new']):
        break
    df['Group'] = df['Group_new']
    df.drop('Group_new', axis=1, inplace=True)

# 循环合并相同Price的组
while True:
    price_min_group = df.groupby('Price')['Group'].min().reset_index()
    df = df.merge(price_min_group, on='Price', suffixes=('', '_new'))
    if df['Group'].equals(df['Group_new']):
        break
    df['Group'] = df['Group_new']
    df.drop('Group_new', axis=1, inplace=True)

# 重新为Group分配连续编号（从0开始）
df['Group'] = df['Group'].rank(method='dense', ascending=True).astype(int) - 1

# 生成目标汇总表
summary_df = df.groupby('Group').agg(
    Price=('Price', list),
    Color=('Color', list)
).reset_index()

为什么单独groupby无法实现？ #

你之前尝试的单一维度groupby，只能基于Color或Price其中一个字段分组，而你的需求是跨维度的连通分组——只要两行在任意一个维度上有共性，就归为同一组，这已经不是简单的单一键分组能解决的，需要用连通分量的思路处理。

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