Pandas DataFrame高级索引：根据Series动态选择特定位置元素

本文旨在探讨如何在pandas dataframe中高效地根据另一个series提供的索引和列名来选择特定位置的元素。针对传统循环方法的性能瓶颈，文章将介绍两种主要的矢量化解决方案：一种结合`factorize`和`reindex`进行2d索引查找，另一种则利用`merge`和`stack`操作实现。这些方法能够显著提升数据处理效率，并适用于需要根据动态条件从dataframe中提取数据的场景。

引言：问题描述与传统方法的局限性

在数据分析实践中，我们经常需要从Pandas DataFrame中提取特定位置的数据。一个常见的场景是，我们拥有一个DataFrame，其行索引和列名均已定义。同时，我们还有一个Pandas Series，该Series的索引与DataFrame的列名相对应，而Series的值则指定了DataFrame中要提取数据的行索引。我们的目标是根据Series提供的这些动态映射关系，从DataFrame中高效地提取相应的元素，并将其组织成一个新的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)
# 输出:
#    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 = pd.Series({'a': 1, 'c': 2, 'b': 3})
print("\nSeries sr:")
print(sr)
# 输出:
# a    1
# c    2
# b    3
# dtype: int64

我们的目标是根据 sr 的指示，提取 df['a'] 的第1行（索引为1）元素，df['c'] 的第2行（索引为2）元素，以及 df['b'] 的第3行（索引为3）元素。

一种直观但效率低下的方法是迭代 sr，逐个查找并提取元素：

result_loop = pd.Series()
for col, row_idx in sr.items():
    result_loop[col] = df.loc[row_idx, col]
print("\n传统循环方法结果:")
print(result_loop)
# 输出:
# a     5
# c    12
# b    16
# dtype: int64

这种循环方式在处理小型数据集时尚可接受，但当DataFrame和Series的规模增大时，其性能会急剧下降，因为它无法利用Pandas底层的矢量化优化。因此，我们需要更高效、矢量化的解决方案。

Face++旷视

Face⁺⁺ AI开放平台

下载

解决方案一：利用 factorize 和 2D 数组索引

Pandas提供了强大的索引机制，我们可以通过将Series的索引和值转换为数值标签，并结合DataFrame的reindex方法，实现高效的2D数组索引查找。这种方法避免了显式循环，充分利用了NumPy的底层优化。

核心思想

1. 将Series的索引（对应DataFrame的列）和值（对应DataFrame的行索引）分别转换为唯一的数值编码和对应的唯一值列表。
2. 使用reindex方法调整DataFrame的行和列，使其与Series中涉及的唯一行索引和列名对齐。
3. 利用NumPy的2D数组索引能力，通过编码后的数值标签直接从重排后的DataFrame的NumPy数组中提取元素。

详细步骤与代码示例

# 1. 对 Series 的值和索引进行 factorize 处理
# a_i 存储 sr 值（行索引）的数值编码，idx 存储 sr 值（行索引）的唯一列表
a_i, idx = pd.factorize(sr)
# a_c 存储 sr 索引（列名）的数值编码，col 存储 sr 索引（列名）的唯一列表
a_c, col = pd.factorize(sr.index)

# 2. 使用 reindex 调整 DataFrame 的行和列顺序，使其与 factorize 结果对齐
# 这一步确保 df 的行和列与 idx 和 col 的顺序一致，方便后续的数值索引
df_reindexed = df.reindex(index=idx, columns=col)

# 3. 将重排后的 DataFrame 转换为 NumPy 数组，并使用数值编码进行 2D 索引
# df_reindexed.to_numpy() 得到一个 NumPy 数组
# a_i 作为行索引，a_c 作为列索引，直接从数组中提取元素
extracted_values = df_reindexed.to_numpy()[a_i, a_c]

# 4. 将提取到的值构建成一个新的 Series，并使用 sr 的原始索引
out_factorize = pd.Series(extracted_values, index=sr.index)

print("\n解决方案一 (factorize + reindex + 2D 索引) 结果:")
print(out_factorize)
# 输出:
# a     5
# c    12
# b    16
# dtype: int64

优点

• 高性能： 利用了Pandas和NumPy底层的矢量化操作，避免了Python循环的开销。
• 简洁高效： 代码逻辑清晰，尤其适用于大规模数据集。

解决方案二：利用 merge 和 stack 操作

另一种实现方式是利用Pandas的stack和merge功能。这种方法通过将DataFrame“扁平化”为Series，然后与Series进行合并，从而实现条件式的数据提取。

核心思想

1. 使用stack()方法将DataFrame转换为一个MultiIndex Series，其中索引包含原始的行索引和列名。
2. 将Series转换为DataFrame，以便可以将其与堆叠后的DataFrame进行合并。
3. 通过merge操作，根据Series的值（作为行索引）和Series的索引（作为列名）来匹配并提取数据。

详细步骤与代码示例

# 1. 将 DataFrame 堆叠 (stack) 为一个 Series，索引为 MultiIndex (行索引, 列名)
# 并将结果 Series 命名为 'out'，方便后续合并
df_stacked = df.stack().rename('out')
print("\nDataFrame df.stack() 结果示例:")
print(df_stacked.head())
# 输出:
# 0  a     0
#    b     1
#    c     2
#    d     3
#    e     4
# dtype: int64

# 2. 将 Series sr 转换为 DataFrame，以便进行合并
# reset_index() 会将 sr 的索引变为一个普通列 ('index')，值变为另一列 (0)
sr_df = sr.reset_index()
print("\nSeries sr.reset_index() 结果:")
print(sr_df)
# 输出:
#   index  0
# 0     a  1
# 1     c  2
# 2     b  3

# 3. 将 sr_df 与 df_stacked 进行合并
# left_on=[0, 'index'] 表示 sr_df 的第0列（sr的值，即行索引）和 'index' 列（sr的索引，即列名）
# right_index=True 表示与 df_stacked 的 MultiIndex 进行匹配
# how='left' 表示保留 sr_df 的所有行
merged_df = sr_df.merge(df_stacked,
                        left_on=[0, 'index'], # 0 是 sr 的值 (行索引), 'index' 是 sr 的索引 (列名)
                        right_index=True,
                        how='left')

# 4. 设置索引并选择结果列
# 将 'index' 列设置回索引，然后选择 'out' 列作为最终结果
out_merge = merged_df.set_index('index')['out']

print("\n解决方案二 (merge + stack) 结果:")
print(out_merge)
# 输出:
# index
# a     5
# c    12
# b    16
# Name: out, dtype: int64

优点

• 可读性： 对于熟悉SQLJOIN操作的用户来说，merge方法可能更易于理解其逻辑。
• 灵活性： merge提供了多种合并方式（left, right, inner, outer），在更复杂的场景下有更大的灵活性。

注意事项与性能考量

1. 性能比较： 通常情况下，factorize结合2D数组索引的方法（解决方案一）在处理大规模数据时会比merge和stack的方法（解决方案二）更快，因为它直接操作NumPy数组，减少了中间数据结构的创建和索引查找的开销。
2. Series索引重复处理：
   • 如果 sr 的索引（即DataFrame的列名）存在重复，例如 sr = pd.Series({'a': 1, 'c': 2, 'a': 3})，则两种方法会有不同的行为。
   • 解决方案一 (factorize): pd.factorize(sr) 默认会保留所有元素，但 reindex 和 2D 索引会根据 factorize 产生的唯一标签进行操作。如果 sr 索引有重复，pd.factorize(sr.index) 会为每个重复项生成不同的编码，导致 out 的索引也会有重复，且值对应 sr 中每个重复项的值。
   • 解决方案二 (merge): merge 操作在遇到左侧（sr_df）有重复键时，会为每个重复键生成一行结果。
   • 建议： 如果 sr 的索引可能存在重复，并且你希望只保留每个列名的一个结果（例如，保留最后一个出现的），则应在执行操作前对 sr 进行去重处理，例如 sr = sr[~sr.index.duplicated(keep='last')]。

总结

本文介绍了两种高效的矢量化方法，用于根据Pandas Series动态选择DataFrame中的特定元素，以替代效率低下的循环操作。第一种方法利用pd.factorize、reindex和NumPy的2D数组索引，在性能上通常表现更优。第二种方法则通过stack和merge操作，提供了另一种逻辑清晰的解决方案。在实际应用中，开发者可以根据数据集规模、性能要求以及个人对代码可读性的偏好，选择最适合的方案。无论选择哪种方法，矢量化操作都是处理Pandas数据以实现高性能计算的关键。