熊猫有条件地创建多列

Question

说我有一个像这样的数据框：

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

我想为每个类别分配不同的颜色值（RGB）。因此，我需要根据z在列class之后插入三列：

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]]

Answer 2

使用merge

df.merge(pd.DataFrame(data=colors,columns=list('rgb')).rename_axis('class').reset_index())
Out[468]: 
   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

Answer 3

join

从colors列表中创建一个数据框，然后加入'class'列。

colors = [[255, 254, 253], [0, 0, 255], [0, 255, 0]]

c = pd.DataFrame(colors, columns=[*'rgb'])
df.join(c, on='class')

   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

---

assign

使用一些数组切片

df.assign(**dict(zip('rgb', np.array(colors)[df['class']].T)))

   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