如何在Python Polars窗口函数中过滤行计算排除式HHI？

针对你的大数据集（近800家企业、204个月、500个市场），最高效的实现方式是通过数学公式转换避免窗口内过滤，直接利用Polars的向量化窗口函数完成计算，完全不需要手动处理。

核心思路 #

排除当前企业的HHI本质是：组内（同一市场+月份）所有其他企业的「规模占排除后总规模的比例」的平方和。

通过数学推导可简化计算：
设：

• ( S ) = 组内所有企业的总规模
• ( s_j ) = 当前企业的规模
• ( \sum s_i^2 ) = 组内所有企业规模的平方和

则排除当前企业后的HHI为：
[
\text{excluded_HHI} = \frac{\sum s_i^2 - s_j^2}{(S - s_j)^2}
]
这个公式将原本O(n²)的计算量压缩到O(n)，完美适配大数据集。

Polars代码实现 #

import polars as pl

# 模拟你的数据集结构（替换为真实数据）
df = pl.DataFrame({
    "market": ["M1"]*3 + ["M2"]*2,
    "month": [202301]*3 + [202301]*2,
    "firm": ["A", "B", "C", "X", "Y"],
    "size": [10, 3, 50, 20, 30]
})

# 计算排除当前企业的HHI
result = df.with_columns(
    # 1. 按市场+月份分组，计算组内总规模、规模平方和
    total_size=pl.col("size").sum().over(["market", "month"]),
    sum_sq=pl.col("size").pow(2).sum().over(["market", "month"])
).with_columns(
    # 2. 用公式计算排除当前企业的HHI
    excluded_hhi=(pl.col("sum_sq") - pl.col("size").pow(2)) / pl.col("total_size").sub(pl.col("size")).pow(2)
).drop("total_size", "sum_sq")  # 可选：删除中间计算列

print(result)

结果验证 #

以示例中2023年企业A为例：

• 组内总规模 ( S = 10+3+50=63 )
• 规模平方和 ( \sum s_i^2 = 10²+3²+50²=2609 )
• 排除A后的HHI = ( (2609-10²)/(63-10)² = 2509/2809 ≈ 0.893 )，与你给出的( (3/53)²+(50/53)² )结果完全一致。

关于窗口过滤的说明 #

如果一定要用窗口过滤的方式实现，Polars也支持，但绝对不推荐用于大数据集——这种方式需要对每个企业遍历组内所有其他企业，时间复杂度为O(n²)，对于800企业的组会产生指数级的计算量，性能极差。示例代码仅作演示：

# 不推荐：窗口过滤实现（仅适合小数据集）
slow_result = df.with_columns(
    excluded_hhi=pl.col("size")
        .filter(pl.col("firm") != pl.col("firm"))
        .map_elements(lambda s: (s / (pl.col("size").sum().over(["market", "month"]) - pl.col("size")))**2)
        .sum()
        .over(["market", "month"])
)

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