Блог

У меня есть рабочий R код для нарезки строк на основе modulus условия, я пытаюсь добиться того же в python любой помощи, которая будет действительно оценена.

R-код:

 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))
 

Ожидаемый результат Python:

 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
 

Давайте попробуем groupby :

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

Обновление: другой подход, который является более подробным, но будет иметь лучшую производительность, особенно для больших данных:

 # 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]
 

Вывод:

     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
 

Попытка избежать функции apply, так что это многословно:

Приведенный ниже код позволяет получить общее количество строк для каждой id группировки, выполнить кумулятивный подсчет и создать контрольный столбец с использованием модуля. Затем запрос отфильтровывает строки на основе вашего условия, прежде чем окончательно сохранить только начальные столбцы исходного фрейма данных.

 (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