Python substitute for below R function to slice rows

user_7120667 0
Viewed 111

I have a working R code to slice rows based on modulus condition, I am trying to achieve same thing in python Any help will be really appreciated.

R Code:

data = read.table(text = 'id  a   b
1   2   0
1   3   0
1   0   0
1   0   1
2   1   0
2   1   1
2   1   0
3   0   1
3   0   1
3   0   1
3   1   0
3   1   1
3   0   1',header=T)

library(dplyr)
data <- data %>%
  group_by(id) %>%
  dplyr::slice(if(n()%% 3 == 0) row_number() else seq_len(n()-1))

Expected Python Output:

id  a   b
1   2   0
1   3   0
1   0   0
2   1   0
2   1   1
2   1   0
3   0   1
3   0   1
3   0   1
3   1   0
3   1   1
3   0   1
pandas dataframe python-3.x numpy python
2 Answers

Let's try groupby:

(df.groupby('id')
   .apply(lambda x: x.iloc[:-1] if len(x)%3 else x)
   .reset_index(level=0,drop=True)
)

Update: Another approach which is more verbose, but will have better performance especially on larger data:

# lazy groupby
g = df.groupby('id')['id']

# number of rows in each group
num_rows=g.transform('size')

# threshold for each group
thresh = np.where(num_rows%3==0, num_rows, num_rows-1)

# enumeration in each group
enum = g.cumcount()

df[enum<thresh]

Output:

    id  a  b
0    1  2  0
1    1  3  0
2    1  0  0
4    2  1  0
5    2  1  1
6    2  1  0
7    3  0  1
8    3  0  1
9    3  0  1
10   3  1  0
11   3  1  1
12   3  0  1
user_42384080

Trying to avoid the apply function, so this is verbose:

What the code below does is get the total number of rows for each id grouping, do a cumulative cumcount, and create a check column using the modulus. The query then filters out rows based on your condition, before finally keeping only the initial columns of the original dataframe.

(df.assign(total=df.groupby("id").id.transform("count"),
           row_number=df.groupby("id").cumcount(),
           filter_check=lambda x: x.total % 3)
  .query("total - row_number != filter_check")
  .filter(df))

    id  a   b
0   1   2   0
1   1   3   0
2   1   0   0
4   2   1   0
5   2   1   1
6   2   1   0
7   3   0   1
8   3   0   1
9   3   0   1
10  3   1   0
11  3   1   1
12  3   0   1
user_71757130
Related