fill values after condition with NaN

Question

I have a df like this:

df = pd.DataFrame(
    [
        ['A', 1],
        ['A', 1],
        ['A', 1],
        ['B', 2],
        ['B', 0],
        ['A', 0],
        ['A', 1],
        ['B', 1],
        ['B', 0]
    ], columns = ['key', 'val'])
df

print:

    key val
0   A   1
1   A   1
2   A   1
3   B   2
4   B   0
5   A   0
6   A   1
7   B   1
8   B   0

I want to fill the rows after 2 in the val column (in the example all values in the val column from row 3 to 8 are replaced with nan).

I tried this:

df['val'] = np.where(df['val'].shift(-1) == 2, np.nan, df['val'])

and iterating over rows like this:

for row in df.iterrows():
    df['val'] = np.where(df['val'].shift(-1) == 2, np.nan, df['val'])

but cant get it to fill nan forward.

Answer 1

You can use boolean indexing with cummax to fill nan values:

df.loc[df['val'].eq(2).cummax(), 'val'] = np.nan

Alternatively you can also use Series.mask:

df['val'] = df['val'].mask(lambda x: x.eq(2).cummax())

---

  key  val
0   A  1.0
1   A  1.0
2   A  1.0
3   B  NaN
4   B  NaN
5   A  NaN
6   A  NaN
7   B  NaN
8   B  NaN

Answer 2

You can try :

ind = df.loc[df['val']==2].index
df.iloc[ind[0]:,1] = np.nan