使用Pandas和NumPy高效地根据条件生成多列

1. 问题背景与传统方法局限性

在数据处理中，我们经常需要根据某些条件从现有数据中派生出新的列。一个常见需求是，从多个潜在的源列中，选择第一个满足条件的非空值作为新列的值，并同时记录该值的原始来源列名。

考虑以下Pandas DataFrame：

import pandas as pd
import numpy as np

data = {'A': [1.0, 2.0, np.nan], 'B': [4, 5, 6]}
df = pd.DataFrame(data)
print("原始DataFrame:")
print(df)

我们期望将其转换为如下形式：

     A  B  val val_source
0  1.0  4  1.0          A
1  2.0  5  2.0          A
2  NaN  6  6.0          B

即，如果列'A'有值，则val取'A'的值，val_source为'A'；否则，val取'B'的值，val_source为'B'。

初学者可能会尝试使用numpy.select，因为它非常适合基于条件进行元素选择。然而，np.select的设计是针对返回单个数组或序列的，它无法直接返回多个独立的列。例如，以下尝试会引发错误：

# 错误的尝试：np.select 无法直接返回多列
# conds = [df['A'].notna(), True]
# choices = [df[['A']].assign(val_source='A'), df[['B']].assign(val_source='B')]
# df[['val', 'val_source']] = np.select(conds, choices) # 这会导致错误

为了解决这个问题，通常需要执行两次独立的np.select操作，分别针对val和val_source列进行赋值，尽管它们的条件是相同的。

# 传统但冗余的方法：两次 np.select 调用
conds = [df['A'].notna(), True]
_choices_val_src = [
    (df['A'], 'A'),
    (df['B'], 'B'),
]
choices_val, choices_src = zip(*_choices_val_src)

df['val'] = np.select(conds, choices_val, default=np.nan)
df['val_source'] = np.select(conds, choices_src, default=np.nan)

print("\n使用两次 np.select 的结果:")
print(df)

这种方法虽然有效，但在代码的简洁性和潜在的性能方面（尤其是在处理大量条件或列时）仍有提升空间。

2. 高效的NumPy索引解决方案

为了实现更简洁、高效的多列条件赋值，我们可以巧妙地结合Pandas和NumPy的特性，特别是NumPy的高级索引功能。核心思想是首先确定每行中哪个源列满足条件（例如，第一个非空值），然后利用这个索引同时提取值和对应的列名。

以下是实现这一目标的优雅方法：

# 重置 DataFrame 以便演示
df = pd.DataFrame(data)

# 1. 确定每行中第一个非NaN值的列索引
# df.notna() 返回一个布尔型DataFrame，指示每个元素是否非空。
# .to_numpy() 将其转换为NumPy数组。
# .argmax(1) 沿着行方向（轴1）找到第一个True（非NaN）值的索引。
# 如果一行中所有值都是NaN，argmax会返回0（即第一列的索引）。
idx = df.notna().to_numpy().argmax(1)

# 2. 使用高级索引提取 'val' 值
# df.to_numpy() 获取DataFrame的底层NumPy数组。
# (df.index, idx) 创建一个元组，其中 df.index 代表行索引，idx 代表每行对应的列索引。
# 这种形式是NumPy高级索引的一种，可以同时选择多个位置的元素。
df['val'] = df.to_numpy()[(df.index, idx)]

# 3. 根据相同的索引提取 'val_source' 列名
# df.columns 是DataFrame的列名索引。
# 使用 idx 直接从 df.columns 中选择对应的列名。
df['val_source'] = df.columns[idx]

print("\n使用NumPy高级索引的最终结果:")
print(df)

代码解析：

1. idx = df.notna().to_numpy().argmax(1):

   • df.notna(): 生成一个与df形状相同的布尔DataFrame，其中True表示非空值，False表示空值。
   • .to_numpy(): 将布尔DataFrame转换为NumPy数组，方便后续的NumPy操作。
   • .argmax(1): 这是关键步骤。它沿着每行（axis=1）查找第一个True值（即第一个非空值）的索引。例如，对于[True, False, True]，argmax(1)会返回0；对于[False, True, True]，会返回1。如果一行中所有值都是False（即所有都是NaN），argmax会返回0，指向第一列。
   • idx现在是一个NumPy数组，包含了每行应该选择哪个列的索引。

2. df['val'] = df.to_numpy()[(df.index, idx)]:

   • df.to_numpy(): 获取原始DataFrame的底层NumPy数组表示。
   • (df.index, idx): 这是一个NumPy的高级索引技巧。df.index提供了行的位置索引（0, 1, 2...），而idx提供了每行对应的列的位置索引。将它们组合成一个元组，NumPy会根据[(row_0_idx, col_0_idx), (row_1_idx, col_1_idx), ...]的形式提取元素。例如，对于第一行，它会提取df.iloc[0, idx[0]]；对于第二行，提取df.iloc[1, idx[1]]，以此类推。这使得我们能够一次性高效地提取所有目标val值。

3. df['val_source'] = df.columns[idx]:

   • df.columns: 这是一个Pandas Index对象，包含了DataFrame的所有列名。
   • df.columns[idx]: 直接使用idx数组作为索引，从df.columns中选取对应的列名。由于idx中的每个元素都代表了该行所选列的索引，因此可以直接映射到列名，从而高效地获取val_source。

3. 注意事项与适用场景

• 效率优势: 这种NumPy高级索引的方法通常比多次使用np.select或循环迭代DataFrame行更高效，尤其是在处理大型数据集时，因为它充分利用了NumPy的矢量化操作。
• 条件限制: 本方案的核心是argmax(1)，它默认查找第一个True值的索引。这意味着它最适用于“选择第一个非空值”或“选择第一个满足特定布尔条件的列”的场景。如果你的条件逻辑更复杂（例如，需要基于数值大小、特定字符串匹配等），并且需要从多个满足条件的列中选择一个，那么可能需要先构建一个更复杂的布尔DataFrame，然后依然可以使用argmax来找到优先级最高的列。
• 全NaN行处理: 如果某一行所有源列都是NaN，df.notna().to_numpy().argmax(1)会返回0（即第一列的索引）。此时，df['val']会取到该行第一列的NaN值，而df['val_source']会是第一列的列名。如果希望全NaN行对应的val和val_source也为NaN，则这种行为是符合预期的。
• 列顺序: 此方法依赖于列的顺序。argmax会选择第一个非空值的列。如果列的顺序很重要（例如，'A'优先于'B'，'B'优先于'C'），请确保DataFrame中的列顺序与你的优先级匹配。

4. 总结

通过巧妙地结合df.notna().to_numpy().argmax(1)来定位目标列索引，并利用NumPy的高级索引能力，我们可以极大地简化和优化在Pandas DataFrame中根据条件生成多列（值及其来源）的代码。这种方法不仅代码量少，而且执行效率高，是处理类似数据转换任务的推荐实践。它展示了理解Pandas底层NumPy操作如何能够帮助我们编写更强大、更高效的数据处理代码。

以上就是使用Pandas和NumPy高效地根据条件生成多列的详细内容，更多请关注php中文网其它相关文章！