Mono Neu Elogy
Mono Neu Elogy

Reputation: 145

Drop rest of data in specific condition in pandas dataframe

I have a dataframe such as:

A       B       C
27.00   9.90    6.24899992
18.00   6.90    4.827007354
15.00   4.20    2.252776065
7.50    2.90    1.673376053
3.00    3.50    3.233439065
4.00    1.20    4.254737365
3.00    2.30    1.257349325
0       8.90    0.254932365
1.00    0.90    2.233293435

now If Column A have 0, I want to drop rest of data in Column B and C from 5 previous row to the bottom row like this:

A       B       C
27.00   9.90    6.24899992
18.00   6.90    4.827007354
15.00   NaN     NaN
7.50    NaN     NaN
3.00    NaN     NaN
4.00    NaN     NaN
3.00    NaN     NaN
0       NaN     NaN
1.00    NaN     NaN

Another example of my dataframe:

A       B       C
27.00   9.90    6.24899992
18.00   6.90    4.827007354
15.00   4.20    2.252776065
7.50    2.90    1.673376053
3.00    NaN     NaN
4.00    NaN     NaN
3.00    NaN     NaN
2.80    NaN     NaN
1.00    NaN     NaN

and The result that I want is just same data because it have not 0 in column A like below:

A       B       C
27.00   9.90    6.24899992
18.00   6.90    4.827007354
15.00   4.20    2.252776065
7.50    2.90    1.673376053
3.00    NaN     NaN
4.00    NaN     NaN
3.00    NaN     NaN
2.80    NaN     NaN
1.00    NaN     NaN

How can I achieve that?

Upvotes: 2

Views: 120

Answers (1)

jezrael
jezrael

Reputation: 862406

If want set 5 values before first 0 and then all values after first 0 to NaNs in columns B and C use:

N = 5
m = df['A'] == 0

idx = next(iter(m.index[m]), df.index[-1] + 1)
print (idx)
2

#if possible less like N rows before first 0 add max
first = max(idx - N, 0)
print (first)
7

df.iloc[first:, df.columns.get_indexer(['B','C'])] = np.nan
print (df)
      A    B         C
0  27.0  9.9  6.249000
1  18.0  6.9  4.827007
2  15.0  NaN       NaN
3   7.5  NaN       NaN
4   3.0  NaN       NaN
5   4.0  NaN       NaN
6   3.0  NaN       NaN
7   0.0  NaN       NaN
8   1.0  NaN       NaN

If no value 0 in column A:

N = 5
m = df['A'] == 0

idx = next(iter(m.index[m]), df.index[-1] + 1)
print (idx)
9

#if possible less like N rows before first 0 add max
first = max(idx - N, 0)
print (first)
4

df.iloc[first:, df.columns.get_indexer(['B','C'])] = np.nan
print (df)
      A    B         C
0  27.0  9.9  6.249000
1  18.0  6.9  4.827007
2  15.0  4.2  2.252776
3   7.5  2.9  1.673376
4   3.0  NaN       NaN
5   4.0  NaN       NaN
6   3.0  NaN       NaN
7  10.0  NaN       NaN
8   1.0  NaN       NaN

First solution:

#create mask
m = df['A'] == 0

#cumulative sum of mask - return Trues for all values after first 0
m1 = m.cumsum() > 0
#counter of values above 0 with swapping order by indexing [::-1] and cumulative sum
s = m.iloc[::-1].cumsum()
#create counter and compare by 5
m2 = s.groupby(s).cumcount() < 5
#chain masks by | for bitwise OR
mask = m1 | m2.sort_index()

#set NaNs by mask
df[['B','C']] = df[['B','C']].mask(mask)
print (df)

Upvotes: 3

Related Questions