Сдвиг строк по индексу с другим значением в кадре данных pandas

Question

Ниже приведены данные выборки

sample_col = ["col1","col2","col3","col4", "col5"]
sample_data = [[1,2,3,4,5],[6,7,8,9,10],[11,12,13,14,15],[16,17,18,19,20]]
sample = pd.DataFrame(data = data_data, columns = data_col)

sample

   col1  col2  col3  col4 col5
 0  1    2     3     4    5
 1  6    7     8     9    10
 2  11   12    13    14   15
 3  16   17    18    19   20

m = [0,1,2,1]

Мне нужно сместить каждую строку в выборке фрейма данных на соответствующее значение в списке м.

Вывод должен выглядеть следующим образомэто

   col1  col2  col3  col4 col5
 0  1    2     3     4    5
 1  NaN  6     7     8    9
 2  NaN  NaN   11    12   13
 3  NaN  16    17    18   19

Я знаю, что могу использовать sample.iloc[i].shift(m[i]) для достижения того, что я хочу, один за другим или все, применяя цикл.Но у меня очень большой набор данных, и на его завершение уходит вечность.

Так есть ли лучший способ достичь того, что я пытаюсь сделать?Любая помощь приветствуется.

Спасибо!

Answer 1

Вы можете использовать логическое срезание, чтобы выбрать только интересующие строки для каждой операции сдвига.Это все равно будет медленным, если set из m большое, но быстрее, чем у вас.

import pandas as pd
import numpy as np

sample_col = ["col1","col2","col3","col4", "col5"]
sample_data = [[1,2,3,4,5],[6,7,8,9,10],[11,12,13,14,15],[16,17,18,19,20]]
sample = pd.DataFrame(data = sample_data, columns = sample_col)
m = [0,1,2,1]

marr = np.array(m)
for i in sorted(set(m)):
    if i==0:
        continue
    sample.loc[marr==i] = sample.loc[marr==i].shift(i, axis=1)