Как найти все строки с одинаковым конкретным значением для столбца, который повторяется определенное время в наборе данных?

Question

У меня есть база данных панд, для которой я хочу найти все строки, имеющие одинаковое значение для столбца A и повторяющиеся определенное количество времени (я называю это size):

   A    B
0  1  yes
1  2   no
2  3   no
3  2  yes
4  3   no
5  4  yes

Так что, если у меня есть size = 2, только значения столбца 2 и 3 повторяются 2 раза в столбце A, поэтому результат должен быть таким:

   A   B1   B2
0  2   no  yes
1  3  yes   no

Я сделал этот код, но онэто довольно медленно для больших данных, так как он использует цикл for, поэтому я ищу предложение для улучшения:

def load_data(self, size):
        '''Creates a data frame of all orders with number of items in that order equals to argument size'''
        data = pd.read_csv(self.file_name, sep="\t", header=0, 

        #grouping
        grouping = data.groupby('A')
        df =  grouping.filter(lambda x: len(x) == size) #filters all the data with specific size
        values = []
        for i in range(0, len(df)):
            if i% size == 0:
                dat = df[i: i+size]
                temp = dat[['B']]
                temp_values = np.append(df.iloc[i][0], temp.values.flatten())
                values.append(temp_values)

        new_df = pd.DataFrame(np.array(values))
        return new_df

Answer 1

g = df.groupby('A')
c = g.cumcount() + 1
s = g.A.transform('size').to_numpy()

df.set_index(['A', c]).B[s == 2].unstack().add_prefix('B').reset_index()

   A  B1   B2
0  2  no  yes
1  3  no   no

---

Если у вас есть больше столбцов

g = df.groupby('A')
c = g.cumcount() + 1
s = g.A.transform('size').to_numpy()


d = df.set_index(['A', c])[s == 2].unstack()
d.columns = [f'{a}{b}' for a, b in d.columns]
d.reset_index()

Answer 2

IIUC, вы можете обойтись с groupby().transform:

df[df.groupby('A').B.transform('size').eq(2)]

дает

    A   B
1   2   no
2   3   no
3   2   yes
4   3   no