本文探讨了在Pandas DataFrame中根据条件从多列中提取值及其来源的常见需求。针对numpy.select无法直接返回多列的限制，文章首先介绍了分离式np.select的传统做法，随后重点阐述并演示了一种利用df.notna().to_numpy().argmax(1)结合高级索引技术，实现更简洁、高效地同时提取目标值和其对应列名的解决方案。
场景描述与问题背景
在数据处理中，我们经常会遇到需要从DataFrame的多个列中，根据特定条件选择一个值，并同时记录该值的来源列名。例如，给定一个包含多列数据的DataFrame，我们可能希望优先选择某个列中的非空值，如果该列为空，则选择下一个备用列中的值，以此类推，并最终生成新的列来存储选定的值及其原始列名。
考虑以下DataFrame作为示例：
import pandas as pdimport numpy as npdata = {'A': [1.0, 2.0, np.nan], 'B': [4, 5, 6]}df = pd.DataFrame(data)print("原始DataFrame:")print(df)
期望的输出是这样的：
A B val val_source0 1.0 4 1.0 A1 2.0 5 2.0 A2 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 indexingdf['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 namesdf['val_source'] = df.columns[idx]print("n使用argmax和高级索引的结果:")print(df)
输出结果：
原始DataFrame: A B0 1.0 41 2.0 52 NaN 6使用两次np.select的结果: A B val val_source0 1.0 4 1.0 A1 2.0 5 2.0 A2 NaN 6 6.0 B使用argmax和高级索引的结果: A B val val_source0 1.0 4 1.0 A1 2.0 5 2.0 A2 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与NumPy：高效地从多列中条件性提取值及来源的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1364176.html
微信扫一扫 
支付宝扫一扫 
