Pandas多列条件提取技巧分享

小伙伴们有没有觉得学习文章很有意思？有意思就对了！今天就给大家带来《Pandas与NumPy：多列条件提取值技巧》，以下内容将会涉及到，若是在学习中对其中部分知识点有疑问，或许看了本文就能帮到你！

场景描述与问题背景

在数据处理中，我们经常会遇到需要从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的列名中提取对应的列名。

下面是实现该方法的步骤和代码：

1. 确定优先级和待处理列： 明确哪些列是我们的目标，并按照优先级顺序排列。在我们的例子中，是A和B。

2. 查找每行中第一个非空值的列索引：

   • 使用df.notna()创建一个布尔型DataFrame，表示每个元素是否非空。
   • 将其转换为NumPy数组to_numpy()。
   • 对每一行（axis=1）使用argmax()。argmax()会返回每行中第一个True值（即第一个非空值）的索引。如果一行中所有值都为False，argmax()会返回0，这需要注意。在我们的场景中，由于True是最后一个条件，它确保了总能找到一个非空值或最后一个备用列。

3. 利用索引提取值：

   • 将DataFrame转换为NumPy数组：df.to_numpy()。
   • 使用高级索引（fancy indexing）来提取值。df.to_numpy()[(df.index, idx)] 允许我们同时指定行索引和列索引来获取特定位置的元素。df.index提供了所有行的索引，而idx提供了每行中我们想要选择的列的索引。

4. 利用索引提取列名：

   • 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中更优雅、高效地处理多列条件选择和来源追踪的问题。

文中关于的知识介绍，希望对你的学习有所帮助！若是受益匪浅，那就动动鼠标收藏这篇《Pandas多列条件提取技巧分享》文章吧，也可关注golang学习网公众号了解相关技术文章。