具有多个条件的 np.where

Question

I Have the following sample dataframe.

     A    B    C    D
1     0    0    0
2     0    0    1
3     1    1    0
4     0    0    1
5    -1    1    1
6     0    0    1
7     0    1    0
8     1    1    1
9     0    0    0
10   -1    0    0

I need to fill column D based on the following rules.

1) I manually assign the first value of D as 0

       df.loc[[0], 'D'] = 0

2) I need to fill the rest of the rows of df['D'] based on the following conditions.

   IF 
        A = 1 AND B = 1 AND D.SHIFT(1) = 0, THEN D = 1
   ELSE IF 
        A = 0, THEN D = D.SHIFT(1)
   ELSE IF 
        A = -1 AND C = 1 AND D.SHIFT(1) = 0, THEN D = -1
   ELSE
        D = 0

The following is my code.

import numpy as np
import pandas as pd

df = pd.DataFrame([[0, 0, 0]]*10, columns=list('ABC'), index=range(1, 11))
df.loc[[3, 8], 'A'] = 1
df.loc[[5, 10], 'A'] = -1
df.loc[[3, 5, 7, 8], 'B'] = 1
df.loc[[2, 4, 5, 6, 8], 'C'] = 1
df.loc[[1], 'D'] = 0
df['D'] = np.where((df.A == 1) & (df.B == 1) & (df.D.shift(1) == 0), 1, 
               np.where((df.A == 0) , df.D.shift(1), 
               np.where((df.A == -1) & (df.C == 1) & (df.D.shift(1) == 0), -1, 0)))

print(df)

The expected output is

   A  B  C    D
1   0  0  0  0
2   0  0  1  0
3   1  1  0  1
4   0  0  1  1
5  -1  1  1  0
6   0  0  1  0
7   0  1  0  0
8   1  1  1  1
9   0  0  0  1
10 -1  0  0  0

But the actual output i receive is the following which is not my desired output.

    A  B  C    D
1   0  0  0  NaN
2   0  0  1  0.0
3   1  1  0  0.0
4   0  0  1  NaN
5  -1  1  1  0.0
6   0  0  1  NaN
7   0  1  0  NaN
8   1  1  1  0.0
9   0  0  0  NaN
10 -1  0  0  0.0

I would really appreciate any help here.

EDIT: the reason I am assigning the first row of df['D'] = 0 is because, df['D'] is also dependent on df['D'].shift(1). If there is any other way to do this even without assigning the first value of df['D'], I am fine with it as well.

Answer 1

Define the following function:

def fn(row):
    if fn.prevD is None:
        fn.prevD = 0
    elif (row.A == 1) & (row.B == 1) & (fn.prevD == 0):
        fn.prevD = 1
    elif row.A == 0:
        pass
    elif (row.A == -1) & (row.C == 1) & (fn.prevD == 0):
        fn.prevD = -1
    else:
        fn.prevD = 0
    return fn.prevD

Then apply it:

fn.prevD = None
df['D'] = df.apply(fn, axis=1)