Pandas，numpy：在`np.select`中返回多列？

5lwkijsr 于 9个月前发布在其他

关注(0)|答案(2)|浏览(115)

我有下面的 Dataframe 。

字符串
我想把它变成下面的框架。

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

型
最简单的方法是什么？我设想一个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语句，尽管它们共享相同的conds。

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)

型
有没有更干净的方法来做这件事？

numpy

来源：https://stackoverflow.com/questions/77544933/pandas-numpy-return-multiple-columns-in-np-select

2条答案

按热度按时间

yuvru6vn1#

使用子集/索引，如下所示：

idx = df.notna().to_numpy().argmax(1)
df['val'] = df.to_numpy()[(df.index, idx)]
df['val_source'] = df.columns[idx]
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

字符串

赞(0）回复(0）举报 9个月前

rkue9o1l2#

对于这组特定的条件，您还可以更改方法：stack和groupby.first。

df[['val_source', 'val']] = df.stack().reset_index(1).groupby(level=0).first()

字符串
输出量：

A  B val_source  val
0  1.0  4          A  1.0
1  2.0  5          A  2.0
2  NaN  6          B  6.0

型
或者因为你只有两个条件可供选择：

m = df['A'].notna()
df['val'] = df['A'].where(m, df['B'])
df['val_source'] = np.where(m, 'A', 'B')

型
输出量：

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

型

赞(0）回复(0）举报 9个月前

我来回答

Pandas，numpy：在`np.select`中返回多列？

2条答案

相关问题

热门标签

最新问答