python - Pandas 条件创建多列

Question

假设我有一个这样的数据框:

x y z class
1 2 3   0
2 2 3   0
1 4 5   2
3 2 2   1

我想为每个类分配不同的颜色值 (RGB)。所以我需要根据 class 在列 z 之后插入三列:

x y z  r   g   b class
1 2 3 255 254 253  0
2 2 3 255 254 253  0
1 4 5  0  255  0   2
3 2 2  0   0  255  1

目前我是这样做的:

# insert three columns
df['r']=0
df['g']=0
df['b']=0
# replace r/g/b values based on `class`
def colorit(dataframe):
    colors = [[255, 254, 253], [0, 0, 255], [0, 255, 0]]
    for i in range(3):
        dataframe.loc[dataframe['c']==i, 'r'] = colors[i][0]
        dataframe.loc[dataframe['c']==i, 'g'] = colors[i][1]
        dataframe.loc[dataframe['c']==i, 'b'] = colors[i][2]

但我认为应该有一些方法可以利用 apply 或 map 方法或类似的方法来更优雅、更高效地完成此操作(使用更少的代码并且没有循环)。

Answer 1

你可以做

In [237]: df.assign(**pd.DataFrame([colors[x] for x in df['class']], columns=['r', 'g', 'b']))
Out[237]:
   x  y  z  class    r    g    b
0  1  2  3      0  255  254  253
1  2  2  3      0  255  254  253
2  1  4  5      2    0  255    0
3  3  2  2      1    0    0  255

---

详细信息

In [238]: df
Out[238]:
   x  y  z  class
0  1  2  3      0
1  2  2  3      0
2  1  4  5      2
3  3  2  2      1

In [239]: colors
Out[239]: [[255, 254, 253], [0, 0, 255], [0, 255, 0]]