Pandas中高效比较两DataFrame值范围并计数匹配项

引言

在数据分析和处理中，我们经常面临需要比较两个dataframe之间数据关系的任务。一个常见的场景是，需要检查一个dataframe（例如df1）中的特定数值列是否落在另一个dataframe（例如df2）定义的数值区间内，并统计每个区间内的匹配数量。虽然通过遍历df1的每一行并对df2进行条件筛选可以实现，但这种基于python循环的迭代方式在处理大数据集时效率低下，严重影响性能。本文将介绍一种利用pandas内置功能，特别是cross merge操作，实现高效、向量化比较的方法。

传统迭代方法的局限性

假设我们有两个DataFrame，df1包含Date、High、Mid、Low等时间序列数据，df2包含Start、Top、Bottom等区间定义。我们的目标是对于df2中的每一行（即每个区间），统计df1中有多少行的High值高于该区间的Bottom，且Mid或Low的最大值低于该区间的Top，同时Date不等于Start。

一个直观但低效的实现方式可能如下所示：

import pandas as pd

# 示例数据
data1 = {
    'Date': pd.to_datetime(['2023-08-03 00:00:00', '2023-08-03 12:00:00', '2023-11-28 00:00:00', '2023-11-29 08:00:00', '2023-11-24 12:00:00']),
    'High': [29249.8, 29395.8, 37341.4, 38419.0, 38432.9],
    'Mid': [29136.6, 29228.1, 37138.2, 38136.3, 37894.4],
    'Low': [29152.3, 29105.0, 37254.1, 38112.0, 37894.4]
}
df1 = pd.DataFrame(data1)

data2 = {
    'Start': pd.to_datetime(['2023-11-28 00:00:00', '2023-11-24 12:00:00']),
    'Top': [37341.4, 38432.9],
    'Bottom': [37138.2, 37894.4]
}
df2 = pd.DataFrame(data2)

df2['Match'] = 0 # 初始化匹配计数

# 迭代方法 (效率低下)
for idx_df1 in df1.index:
    df2.loc[
        (df2.Start != df1.at[idx_df1, 'Date']) &
        (df2.Bottom < df1.at[idx_df1, 'High']) &
        (df2.Top > df1.loc[idx_df1, ['Mid', 'Low']].max()),
        'Match'] += 1

print("迭代方法结果:")
print(df2)

这种方法通过for循环逐行访问df1的数据，并在每次循环中对整个df2进行条件筛选和更新。在Pandas中，直接的Python循环通常比向量化操作慢几个数量级，因此在处理大型数据集时应尽量避免。

高效解决方案：使用 cross merge

为了实现向量化操作，我们可以利用Pandas的merge函数，特别是how='cross'选项，来创建两个DataFrame的笛卡尔积。这样，df1中的每一行都会与df2中的每一行进行配对，形成一个包含所有可能组合的中间DataFrame。在这个大型DataFrame上，我们可以一次性应用所有比较条件，从而实现高效的筛选和计数。

紫东太初

中科院和武汉AI研究院推出的新一代大模型

下载

以下是使用cross merge的实现代码：

import pandas as pd

# 示例数据 (与上面相同)
data1 = {
    'Date': pd.to_datetime(['2023-08-03 00:00:00', '2023-08-03 12:00:00', '2023-11-28 00:00:00', '2023-11-29 08:00:00', '2023-11-24 12:00:00']),
    'High': [29249.8, 29395.8, 37341.4, 38419.0, 38432.9],
    'Mid': [29136.6, 29228.1, 37138.2, 38136.3, 37894.4],
    'Low': [29152.3, 29105.0, 37254.1, 38112.0, 37894.4]
}
df1 = pd.DataFrame(data1)

data2 = {
    'Start': pd.to_datetime(['2023-11-28 00:00:00', '2023-11-24 12:00:00']),
    'Top': [37341.4, 38432.9],
    'Bottom': [37894.4, 37138.2] # 修正df2的Bottom值，使其与df1数据有匹配
}
df2 = pd.DataFrame(data2)

# 为了避免与原始问题答案的输出不一致，这里对df2的Bottom值进行调整，以便产生匹配
# 原始问题答案中的df2:
#                 Start      Top   Bottom
# 0 2023-11-28 00:00:00  37341.4  37138.2
# 1 2023-11-24 12:00:00  38432.9  37894.4
# 示例数据中df2的Bottom值与问题答案中的df2保持一致，以复现其结果
data2_for_answer = {
    'Start': pd.to_datetime(['2023-11-28 00:00:00', '2023-11-24 12:00:00']),
    'Top': [37341.4, 38432.9],
    'Bottom': [37138.2, 37894.4]
}
df2_final = pd.DataFrame(data2_for_answer)


# 使用 cross merge 的高效方法
df2_final['Match'] = (df2_final.reset_index() # 1. 重置 df2 的索引，以便后续 value_counts 能正确引用
                   .merge(df1, how='cross') # 2. 执行 cross merge，生成笛卡尔积
                   .loc[lambda x: # 3. 应用筛选条件
                        (x.Start !=  x.Date) & # 条件1: 日期不相等
                        (x.Bottom < x.High) & # 条件2: df1的High高于df2的Bottom
                        (x.Top > x[['Mid', 'Low']].max(axis=1))] # 条件3: df1的Mid或Low的最大值低于df2的Top
                   .value_counts('index') # 4. 统计每个原始 df2 索引的匹配数量
                   .reindex(df2_final.index, fill_value=0)) # 5. 重新索引 df2，将未匹配的填充为0

print("\ncross merge 方法结果:")
print(df2_final)

代码解析

1. df2_final.reset_index(): 在进行cross merge之前，我们首先对df2_final进行reset_index()操作。这是为了将df2_final原有的索引（在此例中是0和1）转换为一个普通的数据列。这样做的好处是，在cross merge之后，这个原始索引列会被保留下来，方便我们后续使用value_counts()来统计每个原始df2_final行的匹配数量。
2. .merge(df1, how='cross'): 这是核心步骤。how='cross'参数指示Pandas执行两个DataFrame的笛卡尔积。这意味着df1中的每一行都会与df2_final中的每一行进行组合，生成一个拥有len(df1) * len(df2_final)行的新DataFrame。这个新DataFrame包含了所有可能的行组合，使得后续的向量化比较成为可能。
3. .loc[lambda x: ... ]: 在生成笛卡尔积之后，我们使用.loc结合lambda函数来应用所有的筛选条件。x代表合并后的DataFrame。
   • (x.Start != x.Date): 检查合并后行中Start列的值是否不等于Date列的值。
   • (x.Bottom
   • (x.Top > x[['Mid', 'Low']].max(axis=1)) : 检查Top列的值是否大于Mid和Low两列中的最大值。max(axis=1)确保了对每一行进行Mid和Low的比较。 这些条件共同筛选出所有符合匹配逻辑的行。
4. .value_counts('index'): 筛选完成后，我们得到一个只包含匹配行的DataFrame。value_counts('index')的作用是统计原始df2_final的每个索引（在步骤1中被保留为index列）出现了多少次，这直接对应了每个df2_final行有多少个匹配项。
5. .reindex(df2_final.index, fill_value=0): value_counts()只会返回有匹配项的索引及其计数。为了确保df2_final中的所有原始索引都得到一个匹配计数（包括那些没有匹配项的，其计数应为0），我们使用reindex()方法。它将value_counts的结果重新索引到df2_final的原始索引上，对于value_counts中没有出现的索引，fill_value=0会将它们的匹配数设置为0。
6. df2_final['Match'] = ...: 最后，将计算得到的匹配计数结果赋值给df2_final的新列Match。

性能考虑与注意事项

• 内存消耗: cross merge操作会生成两个DataFrame的笛卡尔积，其行数等于len(df1) * len(df2)。如果df1和df2都非常大，这个中间DataFrame可能会消耗大量的内存，甚至导致内存溢出。因此，这种方法最适用于其中一个或两个DataFrame大小适中，能够被内存容纳的场景。
• 计算效率: 尽管存在内存开销，但cross merge结合向量化筛选的计算效率远高于Python循环。Pandas的底层C实现使得大规模的条件判断和聚合操作得以快速执行。
• 替代方案: 对于极大规模的数据集，当cross merge的内存消耗成为瓶颈时，可能需要考虑其他策略，例如：
  • 分块处理 (Chunking): 将一个或两个DataFrame分成小块，分批进行cross merge和处理。
  • 数据库或分布式计算: 利用SQL数据库的JOIN操作或Apache Spark、Dask等分布式计算框架来处理超大数据集。
  • 专门的区间树/KD树: 如果区间查询是核心且非常频繁，可以考虑使用专门的数据结构（如区间树或KD树）来优化查询性能，但这超出了Pandas的范畴。

总结

在Pandas中比较一个DataFrame的数值是否落在另一个DataFrame定义的范围内并计数匹配项是一个常见的需求。通过利用cross merge创建笛卡尔积并结合向量化筛选，可以显著提高处理效率，避免低效的Python迭代。然而，这种方法的适用性受限于内存资源，对于特大数据集需要权衡利弊或考虑其他分布式解决方案。理解并恰当运用cross merge，将有助于数据分析师在处理复杂数据关系时编写出更高效、更专业的代码。