Создать новые столбцы в зависимости от размера элемента groupby в пандах

Question

Вот пример кадра данных,

    id  Section A   B
0   abc foo 0.1 0.6
1   abc foo 0.2 0.3
2   abc bar 0.5 0.1
3   def foo 0.1 0.1
4   def bar 0.1 0.3
5   def bar 0.6 0.1
6   ghj foo 0.3 0.1
7   ghj foo 0.1 0.7
8   ghj bar 0.1 0.2

Новые столбцы df['AA', 'BB'] будут созданы из следующих списков.

A_foo = [0.1,2]
A_bar = [1,0.3]

B_foo = [0.4,0.2]
B_bar = [1.2,0.5]

Вот как я пытался до сих пор,

g = df.groupby('id')['A','B']
for i, i_d in g:
    print(i_d)

**

length of `A_foo, A_bar, B_foo and B_bar` is always greater or equal to df`

[df.Section == 'foo'] and df[df.Section == 'bar']` of any unique id. 

Затем, чтобы создать df['AA'], для каждого 'foo' and 'bar' в df['Section'] для идентификатора я хочу взять соответствующие значения из A_foo and A_bar.

Например, в первом i_d (id = abc) df.A имеет two 'foo' and one 'bar', тогда будут выглядеть первые три строки df['AA'],

[0.1,2,1... 0.1 and 2 from A_foo and 1 from A_bar

затем во втором i_d(id='def'), df.A has one foo and two bar, поэтому мне нужно добавить 0.1 from A_foo and 1,0.3 from A_bar. а теперь

df['AA'] will look [0.1,2,1,0.1,1,0.3...

с последнего i_d, я соберу 0.1,2 from A_foo and 1 from A_bar. теперь полный

df['AA'] = [0.1,2,1,0.1,1,0.3,0.1,2,1]

Аналогично, создайте df['BB'] из B_foo и B_bar

df['BB'] = [0.4,0.2,1.2,0.4,1.2,0.5,0.4,0.2,1.2]

Вот финал df

    id  Section A   B   AA  BB
0   abc foo    0.1  0.6 0.1 0.4
1   abc foo    0.2  0.3 2.0 0.2
2   abc bar    0.5  0.1 1.0 1.2
3   def foo    0.1  0.1 0.1 0.4
4   def bar    0.1  0.3 1.0 1.2
5   def bar    0.6  0.1 0.3 0.5
6   ghj foo    0.3  0.1 0.1 0.4
7   ghj foo    0.1  0.7 2.0 0.2
8   ghj bar    0.1  0.2 1.0 1.2

Answer 1

Создайте индекс с помощью groupby + cumcount, затем используйте np.select для назначения значений из соответствующих списков.

import numpy as np

df['idx'] = df.groupby(['id', 'Section']).cumcount()

conds = [df.Section.eq('foo'), df.Section.eq('bar')]
AA_choice = [np.array(A_foo)[df.idx], np.array(A_bar)[df.idx]]
BB_choice = [np.array(B_foo)[df.idx], np.array(B_bar)[df.idx]]

df['AA'] = np.select(conds, AA_choice, default=np.NaN)
df['BB'] = np.select(conds, BB_choice, default=np.NaN)

Вывод:

    id Section    A    B  idx   AA   BB
0  abc     foo  0.1  0.6    0  0.1  0.4
1  abc     foo  0.2  0.3    1  2.0  0.2
2  abc     bar  0.5  0.1    0  1.0  1.2
3  def     foo  0.1  0.1    0  0.1  0.4
4  def     bar  0.1  0.3    0  1.0  1.2
5  def     bar  0.6  0.1    1  0.3  0.5
6  ghj     foo  0.3  0.1    0  0.1  0.4
7  ghj     foo  0.1  0.7    1  2.0  0.2
8  ghj     bar  0.1  0.2    0  1.0  1.2

---

Если ваши списки недостаточно длинные, вы получите IndexError.Если это так, возможно рассмотрите возможность нарезки на: np.array(A_foo)[df.idx%len(A_foo)]