pandas groupby 中的条件分配

Question

假设我有df以下内容：

df = pd.DataFrame({
    'ID': ['a', 'b', 'b', 'c', 'c', 'c', 'd', 'd', 'd', 'd'],
    'V': np.array(range(0,10))
})

我想要groupby变量ID并将值分配给新列，X具体取决于（1）每个组的大小以及每行是否是顶部（T）、“中间”（更像是顶部和底部之间）（M）或底部（B）排。如果一组中只有一行，则分配的值为N。在这种情况下，结果将如下所示：

    ID  V   X
0   a   0   N
1   b   1   T
2   b   2   B
3   c   3   T
4   c   4   M
5   c   5   B
6   d   6   T
7   d   7   M
8   d   8   M
9   d   9   B

我可以使用类似的方法（对于这种情况）逐步执行此操作T：

df.join(df.groupby('ID').filter(lambda x: len(x)>1).groupby('ID').head(1).assign(X='T').X, how='left')

但这感觉像是一个糟糕的解决方案。我宁愿一次性完成这一切。有任何想法吗？

Answer 1

从你的逻辑来看，这很简单：

groups = df.groupby('ID')
first = groups['V'].head(1).index
last = groups['V'].tail(1).index

# the default middle values
df['X'] = 'M'

# the top and bottom values
df.loc[first, 'X'] = 'T'
df.loc[last, 'X'] = 'B'

# the unique values
ones = groups['V'].transform('size') == 1
df.loc[ones, 'X'] = 'N'

输出：

  ID  V  X
0  a  0  N
1  b  1  T
2  b  2  B
3  c  3  T
4  c  4  M
5  c  5  B
6  d  6  T
7  d  7  M
8  d  8  M
9  d  9  B