求更优方案：计算DataFrame选定列唯一值在其余列组合中的占比

问题 #

我有一个列数可变的DataFrame，需要计算选定列的每个唯一值在其余列的每个唯一组合中的占比。示例如下：

import numpy as np
import pandas as pd

a="A";b="B"
df = pd.DataFrame({
    "n": [a,a,a,b,a,b,b,b],
    "X": [0,0,0,0,1,1,1,1],
    "Y": [0,0,1,1,0,0,0,0],
})
print(df)

输出：

n  X  Y
0  A  0  0
1  A  0  0
2  A  0  1
3  B  0  1
4  A  1  0
5  B  1  0
6  B  1  0
7  B  1  0

以计算唯一值n的绝对频率n_ru在X和Y的每个唯一组合的绝对频率n_u中的占比为例：比如组合(X=1,Y=0)共有4条数据，其中n=B有3条，占比为3/4。

我已实现如下代码：

# complete column list
col = list(df.columns.values)
# column list except n
cov = list(df.columns[1:].values)

# merge absolute frequencies
count = pd.merge(
    # absolute freq of each (X,Y)
    df.groupby(cov).count(),
    # absolute freq of n for each (X,Y)
    df.groupby(col).aggregate("n").count(),
    # options
    on=cov, suffixes=["_u", "_ru"]
)
print(count)
# calculate ell metric
ell = np.sum(
    np.log(count["n_ru"]/count["n_u"])
)
print(f"ell = {ell:.3f}")

输出：

n_ru  n_u
X Y
0 0     2    2
  1     1    2
  1     1    2
1 0     1    4
  0     3    4

ell = -3.060

请问是否有更优的实现方式？

更优实现方案 #

方案一：分组统计+索引对齐（高效简洁） #

通过两次分组统计获取核心数据，用join替代merge（索引对齐更高效），无需手动处理列名列表，逻辑更清晰：

import numpy as np
import pandas as pd

a="A";b="B"
df = pd.DataFrame({
    "n": [a,a,a,b,a,b,b,b],
    "X": [0,0,0,0,1,1,1,1],
    "Y": [0,0,1,1,0,0,0,0],
})

# 指定目标列和分组列（自动适配列数可变场景）
target_col = "n"
group_cols = [col for col in df.columns if col != target_col]

# 1. 计算每个分组（X,Y组合）的总条数
group_totals = df.groupby(group_cols).size().rename("n_u")

# 2. 计算每个分组内目标列各值的出现次数
group_target_counts = df.groupby(group_cols + [target_col]).size().rename("n_ru")

# 3. 按索引合并结果，无需指定匹配列
count = group_target_counts.join(group_totals)

# 计算ell指标
ell = np.sum(np.log(count["n_ru"] / count["n_u"]))
print(count)
print(f"ell = {ell:.3f}")

方案二：transform生成分组总数（步骤更精简） #

通过transform在原DataFrame中直接添加分组总数字段，再一次性聚合得到结果，代码更紧凑：

import numpy as np
import pandas as pd

a="A";b="B"
df = pd.DataFrame({
    "n": [a,a,a,b,a,b,b,b],
    "X": [0,0,0,0,1,1,1,1],
    "Y": [0,0,1,1,0,0,0,0],
})

target_col = "n"
group_cols = [col for col in df.columns if col != target_col]

# 给每行添加对应分组的总条数
df["n_u"] = df.groupby(group_cols)[target_col].transform("count")

# 一次聚合得到分组内目标值计数+分组总数
count = df.groupby(group_cols + [target_col]).agg(
    n_ru=(target_col, "count"),
    n_u=("n_u", "first")  # 同一分组内n_u值相同，取第一个即可
)

ell = np.sum(np.log(count["n_ru"] / count["n_u"]))
print(count)
print(f"ell = {ell:.3f}")

优势说明 #

• 避免了merge操作，减少内存开销与计算步骤，性能更优
• 自动适配列数可变场景，无需手动切片列名，扩展性更强
• 代码逻辑贴合需求：先获取分组总数，再统计分组内目标值数量，可读性更高

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