Вставить строки в пандах, где один столбец пропускает какое-то значение в групповом

Question

Вот мой фрейм данных:

user1    user2    cat    quantity + other quantities
----------------------------------------------------
Alice    Bob      0      ....
Alice    Bob      1      ....
Alice    Bob      2      ....
Alice    Carol    0      ....
Alice    Carol    2      ....

Я хочу убедиться, что каждая пара user1-user2 имеет строку, соответствующую каждой категории (их три: 0,1,2).Если нет, я хочу вставить строку и установить другие столбцы в ноль.

user1    user2    cat    quantity + other quantities
----------------------------------------------------
Alice    Bob      0      ....
Alice    Bob      1      ....
Alice    Bob      2      ....
Alice    Carol    0      ....
Alice    Carol    1      <SET ALL TO ZERO>
Alice    Carol    2      ....

На данный момент у меня есть список всех user1-user2, который имеет менее 3 значений для cat:

df.groupby(['user1','user2']).agg({'cat':'count'}).reset_index()[['user1','user2']]

Я могу перебрать этих пользователей, но это займет много времени (таких пар> 1M).Я проверял другие решения для вставки строк в панды на основе какого-либо условия (например, Pandas / Python добавление строки на основе условия и Вставка строки в Pandas Dataframe на основе условия ), ноони не совсем одинаковые.

Кроме того, поскольку это огромный набор данных, решение должно быть векторизовано.Как мне поступить?

Answer 1

Используйте set_index с reindex от MultiIndex.from_product:

print (df)
   user1  user2  cat  quantity  a
0  Alice    Bob    0         2  4
1  Alice    Bob    1         3  4
2  Alice    Bob    2         4  4
3  Alice  Carol    0         6  4
4  Alice  Carol    2         3  4

df = df.set_index(['user1','user2', 'cat'])
mux = pd.MultiIndex.from_product(df.index.levels, names=df.index.names) 
df = df.reindex(mux, fill_value=0).reset_index()
print (df)
   user1  user2  cat  quantity  a
0  Alice    Bob    0         2  4
1  Alice    Bob    1         3  4
2  Alice    Bob    2         4  4
3  Alice  Carol    0         6  4
4  Alice  Carol    1         0  0
5  Alice  Carol    2         3  4

Другое решение - создание нового Dataframe по всем комбинациям уникальных значений столбцов и merge с right объединение:

from  itertools import product

df1 = pd.DataFrame(list(product(df['user1'].unique(),
                                df['user2'].unique(),
                                df['cat'].unique())), columns=['user1','user2', 'cat'])
df = df.merge(df1, how='right').fillna(0)
print (df)
   user1  user2  cat  quantity    a
0  Alice    Bob    0       2.0  4.0
1  Alice    Bob    1       3.0  4.0
2  Alice    Bob    2       4.0  4.0
3  Alice  Carol    0       6.0  4.0
4  Alice  Carol    2       3.0  4.0
5  Alice  Carol    1       0.0  0.0

EDIT2:

df['user1'] = df['user1'] + '_' + df['user2']
df = df.set_index(['user1', 'cat']).drop('user2', 1)
mux = pd.MultiIndex.from_product(df.index.levels, names=df.index.names)
df = df.reindex(mux, fill_value=0).reset_index()
df[['user1','user2']] = df['user1'].str.split('_', expand=True)
print (df)
   user1  cat  quantity  a  user2
0  Alice    0         2  4    Bob
1  Alice    1         3  4    Bob
2  Alice    2         4  4    Bob
3  Alice    0         6  4  Carol
4  Alice    1         0  0  Carol
5  Alice    2         3  4  Carol

EDIT3:

cols = df.columns.difference(['user1','user2'])
df = (df.groupby(['user1','user2'])[cols]
        .apply(lambda x: x.set_index('cat').reindex(df['cat'].unique(), fill_value=0))
        .reset_index())
print (df)
   user1  user2  cat  a  quantity
0  Alice    Bob    0  4         2
1  Alice    Bob    1  4         3
2  Alice    Bob    2  4         4
3  Alice  Carol    0  4         6
4  Alice  Carol    1  0         0
5  Alice  Carol    2  4         3