Numpy Where 和 Pandas：如何聚合 groupby 值？

Question

如何获得将分组列聚合为单个实体（列表/数组）的数组，同时还为与where子句条件不匹配的结果返回NaNs ？

# example
df1 = pd.DataFrame({'flag': [1, 1, 0, 0],
                    'part': ['a', 'b', np.nan, np.nan],
                    'id': [1, 1, 2, 3]})

# my try
np.where(df1['flag'] == 1, df1.groupby(['id'])['part'].agg(np.array), df1.groupby(['id'])['part'].agg(np.array))
# operands could not be broadcast together with shapes (4,) (3,) (3,) 

# expected
np.array((np.array(('a', 'b')), np.array(('a', 'b')), np.nan, np.nan), dtype=object)

Answer 1

在part列中删除具有NaN值的行，然后按id对剩余的行进行分组，并使用list聚合part ，最后map将聚合的 dataframe 到flag列上以获得结果

s = df1.dropna(subset=['part']).groupby('id')['part'].agg(list)
df1['id'].map(s).to_numpy()

---

array([list(['a', 'b']), list(['a', 'b']), nan, nan], dtype=object)