Панды условного создания мультиколонок

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

Использование 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 2

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

Answer 3

Вы могли бы сделать

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