场景描述与问题背景
在数据处理中，我们经常会遇到需要从DataFrame的多个列中，根据特定条件选择一个值，并同时记录该值的来源列名。例如，给定一个包含多列数据的DataFrame，我们可能希望优先选择某个列中的非空值，如果该列为空，则选择下一个备用列中的值，以此类推，并最终生成新的列来存储选定的值及其原始列名。
考虑以下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
这里，val列存储了根据条件（优先A列，其次B列）选择的值，而val_source列则记录了该值来源于A列还是B列。
numpy.select的局限性
numpy.select是一个非常强大的条件选择函数，它允许我们根据一系列条件和对应的选择值来构建新数组。然而，其设计限制在于它一次只能返回一个数组。当我们需要同时返回多个相关的数组（例如，值和值的来源）时，传统的np.select方法会迫使我们进行多次调用，即使条件集是相同的。
例如，为了实现上述需求，一种常见的做法是执行两次np.select：
conds = [df['A'].notna(), True] # 条件：A不为空，否则为真（即B）
# 准备val的选择项
choices_val = [df['A'], df['B']]
# 准备val_source的选择项
choices_src = ['A', 'B']
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)这种方法虽然功能上可行，但当条件或选择项变得复杂时，代码会显得冗余且可读性下降，因为它重复了条件逻辑。
高效解决方案：利用argmax和高级索引
为了解决np.select的局限性并提供更简洁高效的方案，我们可以利用Pandas和NumPy的强大索引能力。核心思想是：首先找出每行中符合条件的第一个非空值所在的列的索引，然后利用这个索引直接从原始DataFrame的NumPy数组表示中提取值，并从DataFrame的列名中提取对应的列名。
下面是实现该方法的步骤和代码：
确定优先级和待处理列： 明确哪些列是我们的目标，并按照优先级顺序排列。在我们的例子中，是A和B。
-
查找每行中第一个非空值的列索引：
- 使用df.notna()创建一个布尔型DataFrame，表示每个元素是否非空。
- 将其转换为NumPy数组to_numpy()。
- 对每一行（axis=1）使用argmax()。argmax()会返回每行中第一个True值（即第一个非空值）的索引。如果一行中所有值都为False，argmax()会返回0，这需要注意。在我们的场景中，由于True是最后一个条件，它确保了总能找到一个非空值或最后一个备用列。
-
利用索引提取值：
- 将DataFrame转换为NumPy数组：df.to_numpy()。
- 使用高级索引（fancy indexing）来提取值。df.to_numpy()[(df.index, idx)] 允许我们同时指定行索引和列索引来获取特定位置的元素。df.index提供了所有行的索引，而idx提供了每行中我们想要选择的列的索引。
-
利用索引提取列名：
- DataFrame的列名可以通过df.columns获取。
- 直接使用idx来索引df.columns，即可获取对应列的名称。
以下是完整的代码实现：
# 重新初始化df以确保干净状态
data = {'A': [1.0, 2.0, np.nan],
'B': [4, 5, 6]}
df = pd.DataFrame(data)
# 步骤1: 找出每行中第一个非空值的列索引
# df.notna() -> DataFrame of booleans
# .to_numpy() -> Convert to numpy array
# .argmax(1) -> Find the index of the first True (non-NaN) value along axis 1 (rows)
idx = df.notna().to_numpy().argmax(1)
# 步骤2: 利用索引提取值
# df.to_numpy() -> Get the underlying numpy array
# (df.index, idx) -> Create a tuple of row and column indices for advanced indexing
df['val'] = df.to_numpy()[(df.index, idx)]
# 步骤3: 利用索引提取列名
# df.columns -> Get the column names as an Index object
# df.columns[idx] -> Use the calculated column indices to get the corresponding column names
df['val_source'] = df.columns[idx]
print("\n使用argmax和高级索引的结果:")
print(df)输出结果：
原始DataFrame:
A B
0 1.0 4
1 2.0 5
2 NaN 6
使用两次np.select的结果:
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
使用argmax和高级索引的结果:
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注意事项与总结
- 优先级： argmax(1)会返回每行中第一个True的索引。这意味着列的顺序在df.notna()中非常重要。如果你的优先级是A优于B，那么A列必须在B列之前（或者确保df.notna()的列顺序符合你的优先级）。
- 全NaN行处理： 如果某一行所有目标列都为NaN，argmax()会返回0（即第一列的索引）。在这种情况下，df['val']会取第一列的NaN值，而df['val_source']会是第一列的列名。如果需要对这种情况进行特殊处理（例如，设为默认值或另一个来源），则可能需要额外的逻辑。
- 性能： 这种基于NumPy数组操作的方法通常比多次调用Pandas Series操作或Python循环更高效，尤其是在处理大型DataFrame时。
- 可读性： 尽管代码可能看起来有些紧凑，但一旦理解了argmax和高级索引的原理，它比多次np.select更简洁且意图明确。
通过掌握这种利用argmax和高级索引的技巧，可以在Pandas中更优雅、高效地处理多列条件选择和来源追踪的问题。
