本文探讨了在pandas DataFrame中根据条件从多列选择值并追踪其来源的有效方法。针对numpy.select无法直接返回多列的局限性，文章介绍了一种利用DataFrame.notna().argmax(1)结合NumPy高级索引的优化方案，该方案能够简洁高效地实现从多个候选列中提取首个非空值及其对应的列名，显著提升代码的可读性和执行效率。
1. 问题背景：多列条件赋值的挑战
在数据处理中，我们经常需要根据特定条件从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
初次尝试时，开发者可能会倾向于使用numpy.select，因为它允许根据条件数组选择不同的“选择”数组。然而，np.select的一个限制是，它期望每个“选择”项都是一个单一的数组或Series，而不能直接返回一个包含多列的DataFrame。
例如，以下尝试会引发错误：
# 这种尝试会报错，因为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操作，即使它们的条件逻辑完全相同：
# 常见的双次np.select workaround conds = [df['A'].notna(), True] _choices_val_src = [ (df['A'], 'A'), (df['B'], 'B'), ] choices_val, choices_src = zip(*_choices_val_src) df_copy = df.copy() # 使用副本避免修改原始df df_copy['val'] = np.select(conds, choices_val, default=np.nan) df_copy['val_source'] = np.select(conds, choices_src, default=np.nan) print("n使用两次np.select的结果:") print(df_copy)
虽然这种方法能够实现目标，但当需要处理的候选列数量很多时，代码会显得冗长且存在重复计算，不够优雅和高效。
2. 高效解决方案：利用argmax与高级索引
对于“选择第一个非空值及其来源”这类特定场景，Pandas和NumPy提供了一种更为简洁和高效的解决方案，它利用了DataFrame.notna().to_numpy().argmax(axis=1)来找到每行中第一个非空值的列索引，然后结合NumPy的高级索引特性来提取值和列名。
# 原始DataFrame data = {'A': [1.0, 2.0, np.nan], 'B': [4, 5, 6]} df = pd.DataFrame(data) # 1. 识别每行第一个非空值的列索引 # df.notna() 返回一个布尔型DataFrame，指示每个元素是否为非空 # .to_numpy() 将布尔型DataFrame转换为NumPy数组 # .argmax(1) 沿着行方向（axis=1）找到第一个True（即第一个非空值）的索引 # 如果一行全是False（即全是NaN），argmax会返回0（第一个列的索引），这需要注意后续处理 idx = df.notna().to_numpy().argmax(1) print("n每行第一个非空值的列索引 (idx):") print(idx) # 2. 提取对应的值 # df.to_numpy() 将整个DataFrame转换为NumPy数组 # (df.index, idx) 创建一个元组，用于NumPy的高级索引 # df.index 提供了行的索引（0, 1, 2...） # idx 提供了列的索引（0代表A，1代表B...） # 这样，对于每一行，我们都精确地取到了其在idx中指定的列的值 df['val'] = df.to_numpy()[(df.index, idx)] # 3. 提取对应的列名作为来源 # df.columns 是DataFrame的列名列表 # df.columns[idx] 使用之前计算的列索引idx来直接获取对应的列名 df['val_source'] = df.columns[idx] print("n使用argmax和高级索引的结果:") print(df)
代码解析：
-
idx = df.notna().to_numpy().argmax(1):
- df.notna(): 生成一个与df形状相同的布尔型DataFrame，其中True表示非空，False表示空（NaN）。
- .to_numpy(): 将布尔型DataFrame转换为NumPy数组，这是为了利用NumPy的argmax函数。
- .argmax(1): 在NumPy数组上沿着行方向（axis=1）查找第一个True值（即第一个非空值）的索引。如果一行中所有值都为False（即所有列都为NaN），argmax会返回0，表示第一个列的索引。在我们的例子中，由于至少有一个非空值（或者我们希望默认取第一个），这通常是可接受的行为。
-
df[‘val’] = df.to_numpy()[(df.index, idx)]:
- df.to_numpy(): 将整个DataFrame转换为NumPy数组，以便进行高效的NumPy索引操作。
- [(df.index, idx)]: 这是NumPy的高级索引技巧。它接收一个元组，其中第一个元素是行的索引数组，第二个元素是列的索引数组。df.index提供了默认的行索引（0, 1, 2…），而idx提供了我们计算出的每行对应的列索引。通过这种方式，NumPy能够一次性地从df的NumPy表示中精确地提取出每个行在idx指定列上的值。
-
df[‘val_source’] = df.columns[idx]:
- df.columns: 获取DataFrame的所有列名（一个Index对象）。
- df.columns[idx]: 直接使用idx数组作为索引来从df.columns中选择对应的列名。由于idx包含了每个行对应的列索引，这里将直接映射出这些列的名称。
3. 优势与适用场景
- 简洁性: 相较于多次调用np.select，此方法代码量更少，逻辑更集中。
- 效率: 利用了NumPy的矢量化操作，对于大型DataFrame，其执行效率通常高于基于循环或多次条件判断的方法。
- 可读性: 对于“选择第一个有效值”这类特定需求，argmax的语义清晰，代码意图明确。
适用场景：
此方法特别适用于以下场景：
- 从一系列优先级递减的列中选择第一个非空值。
- 每行至少有一个非空值，或者可以接受当所有列都为NaN时默认选择第一个列的值（这可能仍然是NaN）。
注意事项：
- 如果一行中所有列都是NaN，argmax(1)将返回0（即第一列的索引），这意味着val会是第一列的NaN值，而val_source会是第一列的名称。请根据实际业务需求判断这是否是期望的行为。如果需要对这种情况进行特殊处理（例如，赋一个默认值或标记为“无来源”），则可能需要额外的逻辑。
- 此方法侧重于“第一个匹配”的场景。如果条件逻辑更复杂，例如需要基于多个列的组合条件，或者需要选择的不是第一个非空值而是满足特定复杂条件的某个值，那么拆分np.select或使用apply函数可能更具灵活性。
4. 总结
在Pandas中进行多列条件赋值并追踪来源是一个常见的需求。尽管numpy.select在直接返回多列方面存在局限性，但通过巧妙地结合DataFrame.notna().to_numpy().argmax(1)和NumPy的高级索引，我们可以为“选择第一个非空值及其来源”这一特定问题提供一个极其高效和简洁的解决方案。掌握这种技巧将有助于编写更优雅、更具性能的Pandas数据处理代码。
