Nested np.where

后端 未结 3 1123
清酒与你
清酒与你 2021-01-22 03:56

I have the following dataframe:

S A
1 1
1 0
2 1
2 0

I wanted to create a new \'Result\' column that is calculated based on the val

3条回答
  •  深忆病人
    2021-01-22 04:36

    As far as I know np.where does not support multiple return statements (at least not more than two). So either you rewrite your np.where to result in one True and one False statement and to return 1/0 for True/False, or you need to use masks.

    If you rewrite np.where, you are limited to two results and the second result will always be set when the condition is not True. So it will be also set for values like (S == 5) & (A = np.nan).

    df['Result'] = np.where(((df.S == 1) & (df.A == 1)) | ((df.S == 2) & (df.A == 0)), 1, 0)
    

    When using masks, you can apply an arbitrary number of conditions and results. For your example, the solution looks like:

    mask_0 = ((df.S == 1) & (df.A == 0)) | ((df.S == 2) & (df.A == 1))
    mask_1 = ((df.S == 1) & (df.A == 1)) | ((df.S == 2) & (df.A == 0))
    df.loc[mask_0, 'Result'] = 0
    df.loc[mask_1, 'Result'] = 1
    

    Results will be set to np.nan where no condition is met. This is imho failsafe and should thus be used. But if you want to have zeros in these locations, just initialize your Results column with zeros.
    Of course this can be simplified for special cases like only having 1 and 0 as result and extended for any number of result by using dicts or other containers.

提交回复
热议问题