change column in pre-selected elements in pandas dataframe

Viewed 87

We have a dataframe with three different columns, like shown in the example above (df). The goal of this task is to replace the first element of the column 2 by a np.nan, everytime the letter in the column 1 changes. Since the database under study is very big, it cannot be used a for loop. Also every solution that involves a shift is excluded because it is too slow.

I believe the easiest way is to use the groupby and the head method, however I don't know how to replace in the original dataframe.

Examples:

df = pd.DataFrame([['A','Z',1.11],['B','Z',2.1],['C','Z',3.1],['D', 'X', 2.1], ['E','X',4.3],['E', 'X', 2.1], ['F','X',4.3]])

enter image description here

to select the elements that we want to change, we can do the following:

df.groupby(by=1).head(1)[2] = np.nan

However in the original dataframe nothing changes.
The goal is to obtain the following:

enter image description here

Edit:

Based on comments, we won't df[1] returning to a group already seen, e.g. ['Z', 'Z', 'X', 'Z'] is not possible.

5 Answers

Using mask and shift

df[2] = df[2].mask(df[1].ne(df[1].shift(1)))

Using a masked_array:

df[2] = np.ma.masked_array(df[2], df[1].ne(df[1].shift(1))).filled(np.nan)
# array([nan, 2.1, 3.1, nan, 4.3, 2.1, 4.3])

Using np.roll and loc:

a = df[1].values
df.loc[np.roll(a, 1)!=a, 2] = np.nan

   0  1    2
0  A  Z  NaN
1  B  Z  2.1
2  C  Z  3.1
3  D  X  NaN
4  E  X  4.3
5  E  X  2.1
6  F  X  4.3

You can use numpy.where for a vectorised conditional assignment. This assumes df[1] is already sorted.

df[2] = np.where(df[1].duplicated(), df[2], np.nan)

If sorting cannot be assumed:

df[2] = np.where(df[1] != df[1].shift(), np.nan, df[2])

Result:

   0  1    2
0  A  Z  NaN
1  B  Z  2.1
2  C  Z  3.1
3  D  X  NaN
4  E  X  4.3
5  E  X  2.1
6  F  X  4.3

The equivalent logic can be applied using pd.DataFrame.mask.

Using

df[2].mask(df.groupby(1).cumcount().eq(0))
Out[41]: 
0    NaN
1    2.1
2    3.1
3    NaN
4    4.3
5    2.1
6    4.3
Name: 2, dtype: float64

#df[2] = df[2].mask(df.groupby(1).cumcount().eq(0))

Built for speed

a = df[1].values
b = np.flatnonzero(np.append(True, a[1:] != a[:-1]))
df[2].values[b] = np.nan

df

   0  1    2
0  A  Z  NaN
1  B  Z  2.1
2  C  Z  3.1
3  D  X  NaN
4  E  X  4.3
5  E  X  2.1
6  F  X  4.3

You could get the index from your groupby and use it as a mask with .loc.

import pandas as pd
import numpy as np

df = pd.DataFrame([
    ['A','Z',1.11],
    ['B','Z',2.1],
    ['C','Z',3.1],
    ['D', 'X', 2.1], 
    ['E','X',4.3],
    ['E', 'X', 2.1], 
    ['F','X',4.3]
])

m = df.groupby(by=1).head(1).index
df.loc[m,2] = np.nan

print(df)

OR you could use duplicated() instead. This should be faster.

m = ~df[1].duplicated()
df.loc[m, 2] = np.nan

returns:

   0  1    2
0  A  Z  NaN
1  B  Z  2.1
2  C  Z  3.1
3  D  X  NaN
4  E  X  4.3
5  E  X  2.1
6  F  X  4.3
Related