Преобразовать значения столбца в столбец списка на основе условия

Question

Набор данных как:

   id  col2  col3
0   1     1   123
1   1     1   234
2   1     0   345
3   2     1   456
4   2     0  1243
5   2     0   346
6   3     0   888
7   3     0   999
8   3     0   777

Я хотел бы объединить данные по id и добавить значения col3 в список, только если его соответствующее значение в col2 равно 1. Дополнительнодля людей (разных id), у которых только 0 в col2, мне нравится, что агрегированное значение равно 0 для col2 и пустой список для col3.

Вот текущий код:

df_test = pd.DataFrame({'id':[1, 1, 1, 2, 2, 2, 3, 3, 3], 'col2':[1, 1, 0, 1, 0, 0, 0, 0, 0], 'col3':[123, 234, 345, 456, 1243, 346, 888, 999, 777]})

df_test_agg = pd.pivot_table(df_test, index=['id'], values=['col2', 'col3'], aggfunc={'col2':np.max, 'col3':(lambda x:list(x))})

print (df_test_agg)

    col2              col3
id                        
1      1   [123, 234, 345]
2      1  [456, 1243, 346]
3      0   [888, 999, 777]

Желаемый результат должен быть (в идеале, в один шаг в Пандах):

    col2              col3
id                        
1      1            [123, 234]
2      1            [456]
3      0            []

/////////////////////////////////////////////////////////////////////////////////////////

Редактировать - Испытание решения ColdSpeed ​​

df_test = pd.DataFrame({'id':[1, 1, 1, 2, 2, 2, 3, 3, 3], 'col2':[1, 1, 0, 1, 0, 0, 0, 0, 0], 'col3':[123, 234, 345, 456, 1243, 346, 888, 999, 777]})

print (df_test)

df_test_agg = (df_test.where(df_test.col2 > 0)
    .assign(id=df_test.id)
    .groupby('id')
    .agg({'col2': 'max', 'col3': lambda x: x.dropna().tolist()}))

print (df_test_agg)

   id  col2  col3
0   1     1   123
1   1     1   234
2   1     0   345
3   2     1   456
4   2     0  1243
5   2     0   346
6   3     0   888
7   3     0   999
8   3     0   777
    col2            col3
id                      
1    1.0  [123.0, 234.0]
2    1.0         [456.0]
3    NaN              []

///////////////////////////////////////////////////////////////////////////////////////

Отредактировано оригинальное сообщение, чтобы представить больше сценариев.

Answer 1

Вы можете предварительно отфильтровать, а затем использовать groupby:

df_test.query('col2 > 0').groupby('id').agg({'col2': 'max', 'col3': list})

    col2        col3
id                  
1      1  [123, 234]
2      1       [456]

Предостережение заключается в том, что если группа имеет только нули, эта группа будет отсутствовать в результате.Таким образом, чтобы исправить это, вы можете замаскировать с помощью where:

(df_test.where(df_test.col2 > 0)
        .assign(id=df_test.id)
        .groupby('id')
        .agg({'col2': 'max', 'col3'lambda x: x.dropna().tolist()}))

    col2            col3
id                      
1    1.0  [123.0, 234.0]
2    1.0         [456.0]

---

Для обработки 0 групп в "col2" мы можем использовать

(df.assign(col3=df.col3.where(df.col2.astype(bool)))
   .groupby('id')
   .agg({'col2':'max', 'col3': lambda x: x.dropna().astype(int).tolist()}))

    col2        col3
id                  
1      1  [123, 234]
2      1       [456]
3      0          []