本教程探讨了在Pandas DataFrame中根据特定条件（如非空值）高效地创建新列并同时记录其来源的方法。针对numpy.select在直接返回多列时的局限性，文章介绍了一种利用NumPy高级索引和argmax函数实现单行代码解决方案的策略。该方法显著提升了代码的简洁性和执行效率，适用于需要从多个源列中选择第一个有效值及其对应列名的场景。
1. 问题背景与传统方法局限性
在数据处理中，我们经常需要根据某些条件从现有数据中派生出新的列。一个常见需求是，从多个潜在的源列中，选择第一个满足条件的非空值作为新列的值，并同时记录该值的原始来源列名。
考虑以下Pandas 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
即，如果列’A’有值，则val取’A’的值，val_source为’A’；否则，val取’B’的值，val_source为’B’。
初学者可能会尝试使用numpy.select，因为它非常适合基于条件进行元素选择。然而，np.select的设计是针对返回单个数组或序列的，它无法直接返回多个独立的列。例如，以下尝试会引发错误：
# 错误的尝试：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操作，分别针对val和val_source列进行赋值，尽管它们的条件是相同的。
# 传统但冗余的方法：两次 np.select 调用conds = [df['A'].notna(), True]_choices_val_src = [ (df['A'], 'A'), (df['B'], 'B'),]choices_val, choices_src = zip(*_choices_val_src)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)
这种方法虽然有效，但在代码的简洁性和潜在的性能方面（尤其是在处理大量条件或列时）仍有提升空间。
2. 高效的NumPy索引解决方案
为了实现更简洁、高效的多列条件赋值，我们可以巧妙地结合Pandas和NumPy的特性，特别是NumPy的高级索引功能。核心思想是首先确定每行中哪个源列满足条件（例如，第一个非空值），然后利用这个索引同时提取值和对应的列名。
以下是实现这一目标的优雅方法：
# 重置 DataFrame 以便演示df = pd.DataFrame(data)# 1. 确定每行中第一个非NaN值的列索引# df.notna() 返回一个布尔型DataFrame，指示每个元素是否非空。# .to_numpy() 将其转换为NumPy数组。# .argmax(1) 沿着行方向（轴1）找到第一个True（非NaN）值的索引。# 如果一行中所有值都是NaN，argmax会返回0（即第一列的索引）。idx = df.notna().to_numpy().argmax(1)# 2. 使用高级索引提取 'val' 值# df.to_numpy() 获取DataFrame的底层NumPy数组。# (df.index, idx) 创建一个元组，其中 df.index 代表行索引，idx 代表每行对应的列索引。# 这种形式是NumPy高级索引的一种，可以同时选择多个位置的元素。df['val'] = df.to_numpy()[(df.index, idx)]# 3. 根据相同的索引提取 'val_source' 列名# df.columns 是DataFrame的列名索引。# 使用 idx 直接从 df.columns 中选择对应的列名。df['val_source'] = df.columns[idx]print("n使用NumPy高级索引的最终结果:")print(df)
代码解析：
idx = df.notna().to_numpy().argmax(1):
df.notna(): 生成一个与df形状相同的布尔DataFrame，其中True表示非空值，False表示空值。.to_numpy(): 将布尔DataFrame转换为NumPy数组，方便后续的NumPy操作。.argmax(1): 这是关键步骤。它沿着每行（axis=1）查找第一个True值（即第一个非空值）的索引。例如，对于[True, False, True]，argmax(1)会返回0；对于[False, True, True]，会返回1。如果一行中所有值都是False（即所有都是NaN），argmax会返回0，指向第一列。idx现在是一个NumPy数组，包含了每行应该选择哪个列的索引。
df[‘val’] = df.to_numpy()[(df.index, idx)]:
df.to_numpy(): 获取原始DataFrame的底层NumPy数组表示。(df.index, idx): 这是一个NumPy的高级索引技巧。df.index提供了行的位置索引（0, 1, 2…），而idx提供了每行对应的列的位置索引。将它们组合成一个元组，NumPy会根据[(row_0_idx, col_0_idx), (row_1_idx, col_1_idx), …]的形式提取元素。例如，对于第一行，它会提取df.iloc[0, idx[0]]；对于第二行，提取df.iloc[1, idx[1]]，以此类推。这使得我们能够一次性高效地提取所有目标val值。
df[‘val_source’] = df.columns[idx]:
df.columns: 这是一个Pandas Index对象，包含了DataFrame的所有列名。df.columns[idx]: 直接使用idx数组作为索引，从df.columns中选取对应的列名。由于idx中的每个元素都代表了该行所选列的索引，因此可以直接映射到列名，从而高效地获取val_source。
3. 注意事项与适用场景
效率优势: 这种NumPy高级索引的方法通常比多次使用np.select或循环迭代DataFrame行更高效，尤其是在处理大型数据集时，因为它充分利用了NumPy的矢量化操作。条件限制: 本方案的核心是argmax(1)，它默认查找第一个True值的索引。这意味着它最适用于“选择第一个非空值”或“选择第一个满足特定布尔条件的列”的场景。如果你的条件逻辑更复杂（例如，需要基于数值大小、特定字符串匹配等），并且需要从多个满足条件的列中选择一个，那么可能需要先构建一个更复杂的布尔DataFrame，然后依然可以使用argmax来找到优先级最高的列。全NaN行处理: 如果某一行所有源列都是NaN，df.notna().to_numpy().argmax(1)会返回0（即第一列的索引）。此时，df[‘val’]会取到该行第一列的NaN值，而df[‘val_source’]会是第一列的列名。如果希望全NaN行对应的val和val_source也为NaN，则这种行为是符合预期的。列顺序: 此方法依赖于列的顺序。argmax会选择第一个非空值的列。如果列的顺序很重要（例如，’A’优先于’B’，’B’优先于’C’），请确保DataFrame中的列顺序与你的优先级匹配。
4. 总结
通过巧妙地结合df.notna().to_numpy().argmax(1)来定位目标列索引，并利用NumPy的高级索引能力，我们可以极大地简化和优化在Pandas DataFrame中根据条件生成多列（值及其来源）的代码。这种方法不仅代码量少，而且执行效率高，是处理类似数据转换任务的推荐实践。它展示了理解Pandas底层NumPy操作如何能够帮助我们编写更强大、更高效的数据处理代码。
以上就是使用Pandas和NumPy高效地根据条件生成多列的详细内容，更多请关注创想鸟其它相关文章！
版权声明：本文内容由互联网用户自发贡献，该文观点仅代表作者本人。本站仅提供信息存储空间服务，不拥有所有权，不承担相关法律责任。
如发现本站有涉嫌抄袭侵权/违法违规的内容， 请发送邮件至 chuangxiangniao@163.com 举报，一经查实，本站将立刻删除。
发布者：程序猿，转转请注明出处：https://www.chuangxiangniao.com/p/1364174.html
微信扫一扫 
支付宝扫一扫 
