Найти связанное значение в столбце фрейма данных

Question

У меня есть фрейм данных с идентификаторами, именами и флагами.

Я хочу создать новый столбец, в котором будут заполнены связанные идентификаторы.

Вот некоторые правила:

• Если flag = '0' или 'A', то нет связанного идентификатора
• Если flag = 'B', нам нужно найти то же имя (минус флаг в имени ), но с флагом 'A' и получить идентификатор
• Если флаг = 'C', нам нужно найти то же имя (без флага в имени), но с флагом 'B' и получить идентификатор

Вот код для DF:

import pandas as pd

d = {'name': ['test1 A', 'test1 B', 'test2 A', 'test2 B', 'test3','test4 A','test4 B','test4 C'],
     'id': [1, 2, 3, 4, 5, 6, 7, 8],
     'flag': ['A', 'B', 'A', 'B', '0', 'A', 'B', 'C']}
df = pd.DataFrame(data=d)

И это желаемый результат:

name        id     flag    related id
test1 A      1        A           nan
test1 B      2        B           1
test2 A      3        A           nan
test2 B      4        B           3
test3 0      5        0           nan
test4 A      6        A           nan
test4 B      7        B           6
test4 C      8        C           7

Я думал об удалении флага из имя, а затем каким-то образом использовать groupby или merge, но не удалось.

Answer 1

groupby с shift

df['related id']=df.groupby(df.name.str.split().str[0]).id.shift()
df
Out[11]: 
      name  id flag  related id
0  test1 A   1    A         NaN
1  test1 B   2    B         1.0
2  test2 A   3    A         NaN
3  test2 B   4    B         3.0
4    test3   5    0         NaN
5  test4 A   6    A         NaN
6  test4 B   7    B         6.0
7  test4 C   8    C         7.0