高效选取Pandas DataFrame指定元素教程

本教程详细介绍了如何高效地从pandas dataframe中根据一个series的索引和值选取特定元素。文章对比了传统循环方法的低效性，并提供了两种高性能的向量化解决方案：一种利用`pd.factorize`、`reindex`和numpy高级索引，另一种则采用`merge`操作，旨在帮助用户提升数据处理效率。

在数据分析和处理中，我们经常需要从Pandas DataFrame中根据特定的条件选取数据。一个常见的场景是，我们有一个DataFrame df，以及一个Series sr。这个sr的索引对应着df的列名，而sr的值则对应着df的行索引。我们的目标是根据sr中定义的这些“行索引-列名”对，从df中高效地提取相应的数值，并将其组织成一个新的Series或列表。

问题描述与传统方法

假设我们有如下的DataFrame df 和 Series sr：

import pandas as pd
import numpy as np

# 示例 DataFrame
data = np.arange(25).reshape(5, 5)
df = pd.DataFrame(data, columns=list('abcde'))
print("DataFrame df:")
print(df)

# 示例 Series
# sr 的索引是 df 的列名，sr 的值是 df 的行索引
sr = pd.Series([1, 2, 3], index=['a', 'c', 'b'])
print("\nSeries sr:")
print(sr)

输出：

DataFrame df:
    a   b   c   d   e
0   0   1   2   3   4
1   5   6   7   8   9
2  10  11  12  13  14
3  15  16  17  18  19
4  20  21  22  23  24

Series sr:
a    1
c    2
b    3
dtype: int64

我们的目标是根据sr中的信息，获取df.loc[sr.loc['a'], 'a'] (即 df.loc[1, 'a'] -> 5)，df.loc[sr.loc['c'], 'c'] (即 df.loc[2, 'c'] -> 12)，以及df.loc[sr.loc['b'], 'b'] (即 df.loc[3, 'b'] -> 16)。

一个直观但效率较低的方法是遍历sr，并逐个查找元素：

result_loop = pd.Series()
for col_label, row_label in sr.items():
    result_loop[col_label] = df.loc[row_label, col_label]
print("\n结果 (循环方法):")
print(result_loop)

输出：

结果 (循环方法):
a     5
c    12
b    16
dtype: int64

这种方法虽然能得到正确结果，但由于涉及到Python层面的循环，对于大型数据集来说性能会非常差，因为它没有充分利用Pandas和NumPy的向量化操作优势。

优化方案一：利用 pd.factorize、reindex 和 NumPy 高级索引

这种方法通过将标签（行索引和列名）转换为整数位置，然后利用NumPy的高级索引功能进行快速查找，从而实现高效的向量化操作。

Civitai

AI艺术分享平台！海量SD资源和开源模型。

下载

# 1. 对 sr 的值（行索引）和 sr 的索引（列名）进行因子化，获取其整数编码和唯一值
row_codes, unique_rows = pd.factorize(sr)
col_codes, unique_cols = pd.factorize(sr.index)

# 2. 对 df 进行 reindex，使其行索引和列名与 sr 中涉及的唯一值对齐
# 这一步确保 df 的内部顺序与 factorize 得到的 unique_rows/unique_cols 保持一致
# 从而使得后续的 NumPy 整数索引能够准确映射
reindexed_df = df.reindex(index=unique_rows, columns=unique_cols)

# 3. 将对齐后的 DataFrame 转换为 NumPy 数组，并使用整数编码进行高级索引
# row_codes 提供了在 unique_rows 中的位置
# col_codes 提供了在 unique_cols 中的位置
extracted_values = reindexed_df.to_numpy()[row_codes, col_codes]

# 4. 将结果封装回 Series，并使用 sr 的原始索引
result_factorize = pd.Series(extracted_values, index=sr.index)
print("\n结果 (factorize + reindex + NumPy 索引):")
print(result_factorize)

原理分析：

1. pd.factorize(sr) 和 pd.factorize(sr.index)：将sr中的行索引标签和列名标签分别转换为从0开始的整数编码（row_codes, col_codes），同时返回这些标签的唯一值列表（unique_rows, unique_cols）。
2. df.reindex(index=unique_rows, columns=unique_cols)：这一步是关键。它根据sr中涉及到的唯一行索引和列名，重新构建了一个DataFrame。这样做的目的是确保reindexed_df的行和列的内部顺序与unique_rows和unique_cols列表的顺序完全一致。
3. reindexed_df.to_numpy()[row_codes, col_codes]：将reindexed_df转换为NumPy数组后，row_codes和col_codes可以直接作为NumPy的高级索引，分别指定要提取的行位置和列位置。例如，如果row_codes是[0, 1, 2]而col_codes是[0, 1, 2]，则会提取array[0,0], array[1,1], array[2,2]处的元素。由于reindexed_df的行/列已根据unique_rows/unique_cols对齐，因此row_codes和col_codes能准确指向原始df中期望的元素。

这种方法在处理大量数据时表现出卓越的性能，是推荐的首选方案。

优化方案二：利用 merge 操作

另一种向量化的方法是利用Pandas的merge功能。它通过将sr和df进行重塑，然后像数据库连接一样进行匹配。

# 1. 将 sr 转换为 DataFrame，使其行索引和值成为两列
# 'index' 列包含原始 sr 的索引（即 df 的列名）
# 0 列包含原始 sr 的值（即 df 的行索引）
sr_df = sr.reset_index()

# 2. 将 df 堆叠（stack）成 Series，并重命名
# 结果是一个 MultiIndex Series，索引为 (行索引, 列名)
# 值为 df 对应位置的元素
df_stacked = df.stack().rename('out')

# 3. 执行合并操作
# left_on=[0, 'index']：将 sr_df 的第 0 列（行索引）和 'index' 列（列名）作为合并键
# right_index=True：将 df_stacked 的 MultiIndex 作为右侧的合并键
# how='left'：保留 sr_df 中的所有项，并匹配 df_stacked 中的值
merged_result = sr_df.merge(df_stacked,
                            left_on=[0, 'index'],
                            right_index=True,
                            how='left')

# 4. 设置回原始 sr 的索引，并选取结果列
result_merge = merged_result.set_index('index')['out']
print("\n结果 (merge 方法):")
print(result_merge)

原理分析：

1. sr.reset_index()：将sr转换为一个两列的DataFrame，其中一列是原sr的索引（即df的列名），另一列是原sr的值（即df的行索引）。
2. df.stack().rename('out')：将df从宽格式转换为长格式。结果是一个Series，其索引是一个MultiIndex，由(行索引, 列名)组成，值为df中对应位置的元素。
3. merge(...)：执行左连接。left_on=[0, 'index']指定了sr_df中用于匹配的列（分别是sr的值和sr的索引），right_index=True指定了df_stacked的MultiIndex作为匹配键。这样，merge操作会根据sr中定义的(行索引, 列名)对，从df_stacked中查找并提取对应的值。
4. .set_index('index')['out']：将结果DataFrame的索引设置回原始sr的索引（即列名），并选择包含提取值的'out'列。

这种方法在逻辑上更接近于SQL的JOIN操作，对于熟悉数据库操作的用户来说可能更容易理解。

注意事项

• sr 中索引重复的处理：如果sr的索引（即df的列名）存在重复，例如sr = pd.Series([1, 2, 3, 4], index=['a', 'c', 'b', 'a'])，上述两种优化方法都会保留sr中所有项，并为每个重复项生成对应的结果。如果需要对重复项进行特定的处理（例如只保留第一个或最后一个，或进行聚合），应在应用这些方法之前对sr进行预处理，例如：

  # 示例：如果sr中有重复索引，保留最后一个
  # sr_deduplicated = sr[~sr.index.duplicated(keep='last')]
  # 然后将 sr_deduplicated 作为输入

• 性能考量：通常情况下，基于factorize和NumPy高级索引的方法在纯粹的查找性能上会略优于merge方法，尤其是在数据集非常庞大时。merge方法涉及到更多的数据重塑和哈希表操作。
• 可读性与复杂性：merge方法可能在概念上更直观，因为它模拟了常见的数据库连接。而factorize方法虽然性能更优，但其内部机制（标签到整数编码的转换）可能需要一些时间来理解。

总结

本文介绍了从Pandas DataFrame中根据Series的索引和值高效选取元素的两种向量化方法。相比于低效的循环遍历，无论是利用pd.factorize结合reindex和NumPy高级索引，还是采用merge操作，都能显著提升数据处理的效率。在实际应用中，应根据数据规模、性能要求以及个人对不同方法的熟悉程度来选择最合适的方案。