如何在 pandas 的 np.where 子句上使用字典答案

【问题标题】：How to use dictionary on np.where clause in pandas如何在 pandas 的 np.where 子句上使用字典
【发布时间】：2023-02-16 18:41:49
【问题描述】：

我有以下数据框

import pandas as pd
foo = pd.DataFrame({'id': [1,1,1,2,2,2],
                    'time': [1,2,3,1,2,3],
             'col_id': ['ffp','ffp','ffp', 'hie', 'hie', 'ttt'],
             'col_a': [1,2,3,4,5,6],
             'col_b': [-1,-2,-3,-4,-5,-6],
                'col_c': [10,20,30,40,50,60]})

id  time col_id  col_a  col_b  col_c
0   1     1    ffp      1     -1     10
1   1     2    ffp      2     -2     20
2   1     3    ffp      3     -3     30
3   2     1    hie      4     -4     40
4   2     2    hie      5     -5     50
5   2     3    ttt      6     -6     60

我想在foo 中创建一个新的col，它将采用col_a 或col_b 或col_c 的值，具体取决于col_id 的值。

我正在做以下事情：

foo['col'] = np.where(foo.col_id == "ffp", foo.col_a, 
                      np.where(foo.col_id == "hie",foo.col_b, foo.col_c))

这使

  id  time col_id  col_a  col_b  col_c  col
0   1     1    ffp      1     -1     10    1
1   1     2    ffp      2     -2     20    2
2   1     3    ffp      3     -3     30    3
3   2     1    hie      4     -4     40   -4
4   2     2    hie      5     -5     50   -5
5   2     3    ttt      6     -6     60   60

由于我有很多专栏，我想知道是否有更简洁的方法来做到这一点，例如使用字典：

dict_cols_matching = {"ffp" : "col_a", "hie": "col_b", "ttt": "col_c"}

有任何想法吗？

【问题讨论】：

标签： python pandas

【解决方案1】：

你可以mapcol_id 的值与字典，然后执行索引查找：

idx, cols = pd.factorize(foo['col_id'].map(dict_cols_matching))

foo['col'] = foo.reindex(cols, axis=1).to_numpy()[np.arange(len(foo)), idx]

输出：

   id  time col_id  col_a  col_b  col_c  col
0   1     1    ffp      1     -1     10    1
1   1     2    ffp      2     -2     20    2
2   1     3    ffp      3     -3     30    3
3   2     1    hie      4     -4     40   -4
4   2     2    hie      5     -5     50   -5
5   2     3    ttt      6     -6     60   60

【讨论】：