Pandas DataFrame NumPy数组高效查找方法

在数据分析和处理中，我们经常需要根据一组特定的索引或条件从Pandas DataFrame中提取数据。一种常见的、但效率低下的做法是使用for循环逐个查找。当数据集规模较小时，这种方法可能不明显，但对于大型数据集，其性能瓶颈会非常突出。

低效的循环查找示例

假设我们有一个Pandas DataFrame df，其索引包含一系列数值，并且我们希望根据一个NumPy数组 ex_arr 中指定的索引值，批量获取HHt列对应的数据。

import pandas as pd
import numpy as np

# 示例 DataFrame
data = {'HHt': [1, 2, 3, 4, 5, 6, 7, 8]}
index_values = [2643, 2644, 2645, 2646, 2647, 2648, 2649, 2650]
df = pd.DataFrame(data, index=index_values)

# 待查找的索引数组
ex_arr = np.array([2643, 2644, 2647])

print("原始 DataFrame:")
print(df)
print("\n待查找索引数组:", ex_arr)

# 低效的 for 循环方式
print("\n使用 for 循环查找 (低效):")
result_for_loop = []
for i in ex_arr:
    h_p = df.at[i, "HHt"] # df.at 用于单个标量快速访问
    result_for_loop.append(h_p)
print(result_for_loop)

上述代码虽然能得到正确结果 [1, 2, 5]，但其逐个访问的特性在处理大量数据时会导致显著的性能开销。

向量化查找：使用 DataFrame.loc

Pandas 提供了强大的向量化操作，可以极大地提高数据处理效率。对于基于标签（索引值或列名）的批量查找，DataFrame.loc 是理想的选择。它可以接受一个列表或NumPy数组作为行索引，实现一次性查找多个行。

要实现上述需求的向量化查找，我们只需将 ex_arr 直接传递给 df.loc 的行索引部分，并指定需要获取的列名。

# 向量化查找：使用 df.loc
print("\n使用 df.loc 向量化查找:")
# 直接通过 df.loc[行索引, 列索引] 获取数据
# 结果是一个 Pandas Series
result_series = df.loc[ex_arr, 'HHt']
print("结果 (Pandas Series):")
print(result_series)

输出的 result_series 是一个Pandas Series，其中包含我们所需的数据，并且保留了原始索引。

结果 (Pandas Series):
2643    1
2644    2
2647    5
Name: HHt, dtype: int64

将结果转换为列表或NumPy数组

在某些场景下，我们可能需要将查找结果转换为标准的Python列表或NumPy数组。Pandas Series 对象提供了方便的方法来实现这一点：to_list() 和 to_numpy()。

# 将结果转换为 Python 列表
result_list = df.loc[ex_arr, 'HHt'].to_list()
print("\n结果 (Python 列表):")
print(result_list)

# 将结果转换为 NumPy 数组
result_numpy = df.loc[ex_arr, 'HHt'].to_numpy()
print("\n结果 (NumPy 数组):")
print(result_numpy)

这将分别输出：

结果 (Python 列表):
[1, 2, 5]

结果 (NumPy 数组):
[1 2 5]

注意事项

1. loc 与 iloc 的区别：
   • df.loc 是基于标签（索引值和列名）进行选择的。
   • df.iloc 是基于整数位置进行选择的。确保根据你的索引类型选择正确的方法。本教程中，ex_arr 包含的是DataFrame的索引标签，因此 loc 是正确的选择。
2. 索引存在性：如果 ex_arr 中的某个索引值在DataFrame的索引中不存在，df.loc 会抛出 KeyError。在实际应用中，如果存在这种可能性，你可能需要先检查索引是否存在，或者使用 reindex 等方法进行更复杂的处理。
3. 性能优势：向量化操作的性能优势在处理大数据集时尤为显著。Pandas和NumPy的底层实现是C语言，能够高效地执行批处理操作，远超Python层面的for循环。
4. 多列选择：如果需要获取多列数据，可以将列名以列表形式传递给 loc：df.loc[ex_arr, ['ColumnA', 'ColumnB']]。

总结

通过本教程，我们了解了如何利用 DataFrame.loc 方法结合NumPy数组，在Pandas DataFrame中高效地执行向量化查找。这种方法不仅显著提升了数据处理性能，也使得代码更加简洁和可读。掌握向量化操作是高效使用Pandas进行数据分析的关键技能之一。