如何在Pandas DataFrame中匹配行列值并生成结果列d？

问题描述 #

现有如下结构的Pandas DataFrame：

| a| b_1       | b_2       |b_3  |c_1        | c_2       | c_3 |
|--|-----------|-----------|-----|-----------|-----------|-----|
|e1|3295.000000|-775.000000|604.5|3575.000000|-626.000000|604.5|
|e2|3615.000000|-731.000000|604.5|1          |0          |0    |
|e3|3615.000000|-731.000000|604.5|3575.000000|-626.000000|604.5| 
|e2|3615.000000|-731.000000|604.5|1          |0          | 0   | 
|e1|3295.000000|-775.000000|604.5|3575.000000|-626.000000|604.5|
|e4| 0         |0          | 0   |0          |0          | 0   |

需要转换为如下格式的结果DataFrame（列b、c存储3维numpy数组，列d按规则生成）：

|a | b                              |c                              | d |  
|--|--------------------------------|-------------------------------|------|
|e1| [3295.000000,-775.000000,604.5]|[3575.000000,-626.000000,604.5]|e3| 
|e2| [3615.000000,-731.000000,604.5]|[1, 0, 0]                      |e3| 
|e3| [3615.000000,-731.000000,604.5]|[3575.000000,-626.000000,604.5]|e1, e2|
|e4| [0, 0, 0]                      |[0, 0, 0]                      |None  |

列d的生成规则：

1. 若当前行的b值与其他行（排除自身）的b值相等，取这些行的a值；
2. 若当前行的b值与其他行（排除自身）的c值相等，取这些行的a值；
3. 若当前行的c值与其他行（排除自身）的c值相等，取这些行的a值；
4. 无匹配项则为None。

实现步骤与代码 #

1. 导入依赖并构造原始数据 #

import pandas as pd
import numpy as np

# 构造原始DataFrame
data = {
    'a': ['e1', 'e2', 'e3', 'e2', 'e1', 'e4'],
    'b_1': [3295.0, 3615.0, 3615.0, 3615.0, 3295.0, 0.0],
    'b_2': [-775.0, -731.0, -731.0, -731.0, -775.0, 0.0],
    'b_3': [604.5, 604.5, 604.5, 604.5, 604.5, 0.0],
    'c_1': [3575.0, 1.0, 3575.0, 1.0, 3575.0, 0.0],
    'c_2': [-626.0, 0.0, -626.0, 0.0, -626.0, 0.0],
    'c_3': [604.5, 0.0, 604.5, 0.0, 604.5, 0.0]
}
df = pd.DataFrame(data)

2. 合并列并去重 #

将b_1/b_2/b_3、c_1/c_2/c_3分别合并为numpy数组列，同时按a列去重（原始数据存在重复行）：

# 自定义函数：将列转为numpy数组，若全为整数则转int类型
def to_np_array(row, cols):
    arr = row[cols].values
    if np.all(arr == arr.astype(int)):
        return arr.astype(int)
    return arr

# 生成b、c列
df['b'] = df.apply(lambda x: to_np_array(x, ['b_1', 'b_2', 'b_3']), axis=1)
df['c'] = df.apply(lambda x: to_np_array(x, ['c_1', 'c_2', 'c_3']), axis=1)

# 按a列去重，保留每个a的唯一行
df_unique = df.drop_duplicates(subset='a').reset_index(drop=True)

3. 构建数组与a值的映射表 #

由于numpy数组不可哈希，将其转为元组作为字典key，收集所有b、c数组对应的a值：

# 构建b数组到a值的映射
b_map = {}
for _, row in df_unique.iterrows():
    key = tuple(row['b'])
    if key not in b_map:
        b_map[key] = set()
    b_map[key].add(row['a'])

# 构建c数组到a值的映射
c_map = {}
for _, row in df_unique.iterrows():
    key = tuple(row['c'])
    if key not in c_map:
        c_map[key] = set()
    c_map[key].add(row['a'])

4. 生成列d #

根据规则匹配符合条件的a值，合并为字符串或设为None：

def calculate_d(row):
    current_a = row['a']
    current_b = tuple(row['b'])
    current_c = tuple(row['c'])
    matched = set()

    # 规则1：匹配其他行的b值
    if current_b in b_map:
        matched.update(b_map[current_b] - {current_a})
    # 规则2：匹配其他行的c值（用当前b值）
    if current_b in c_map:
        matched.update(c_map[current_b] - {current_a})
    # 规则3：匹配其他行的c值（用当前c值）
    if current_c in c_map:
        matched.update(c_map[current_c] - {current_a})

    return ', '.join(sorted(matched)) if matched else None

# 应用函数生成d列
df_unique['d'] = df_unique.apply(calculate_d, axis=1)

5. 整理最终结果 #

保留需要的列并输出：

result_df = df_unique[['a', 'b', 'c', 'd']]
print(result_df)

输出结果与目标格式一致：

a                              b                              c         d
0  e1  [3295.0, -775.0, 604.5]  [3575.0, -626.0, 604.5]        e3
1  e2  [3615.0, -731.0, 604.5]           [1, 0, 0]        e3
2  e3  [3615.0, -731.0, 604.5]  [3575.0, -626.0, 604.5]  e1, e2
3  e4           [0, 0, 0]           [0, 0, 0]     None

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