Pandas与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)

代码解析：

1. 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，表示第一个列的索引。在我们的例子中，由于至少有一个非空值（或者我们希望默认取第一个），这通常是可接受的行为。

2. 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指定列上的值。

3. 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数据处理代码。

以上就是Pandas与NumPy：高效实现多列条件赋值与来源追踪的详细内容，更多请关注php中文网其它相关文章！