python - Groupby 一列，然后比较其他两列并在不同列中返回一个值

Question

我有一个类似的数据框

    data={'COMB':["PNR1", "PNR1", "PNR11", "PNR2", "PNR2"],
        'FROM':["MAA", "BLR", "DEL", "TRV", "HYD"],
         'TO':["BLR", "MAA", "MAA", "HYD", "TRV"]}
md=pd.DataFrame(data)
md

我想要做的是根据条件创建另一列，如果一行的 From 等于下一行的 To，那么它应该返回“R”，否则它将在新列中返回“O”。我的最终输出应该是这样的。

任何人都可以在 python 中帮助我。我尝试了以下方法，但它给了我错误

md_merged=(md>>
            group_by('COMB')>>
            mutate(TYPE=np.where(md['FROM'].isin(md['TO']),"R","O"))>>
           ungroup)

ValueError：值的长度与索引的长度不匹配请帮助。

Answer 1

此解决方案比较组之间的所有值，不仅是前一个值和下一个值。

GroupBy.apply您可以在布尔掩码中使用自定义 lambda 函数，以避免MultiIndex添加group_keys=False到DataFrame.groupby，最后在 中设置新值numpy.where：

mask = md.groupby('COMB', group_keys=False).apply(lambda x: x['FROM'].isin(x['TO']))
md = md.assign(Type=np.where(mask,"R","O"))
print (md)
    COMB FROM   TO Type
0   PNR1  MAA  BLR    R
1   PNR1  BLR  MAA    R
2  PNR11  DEL  MAA    O
3   PNR2  TRV  HYD    R
4   PNR2  HYD  TRV    R

---

此解决方案比较每组的上一行和下一行：

另一个想法是使用DataFrameGroupBy.shift，它应该更快，例如groupby.apply：

mask = (md.groupby('COMB')['FROM'].shift().eq(md['TO']) | 
        md.groupby('COMB')['TO'].shift(-1).eq(md['FROM']))

md = md.assign(Type=np.where(mask,"R","O"))
print (md)
    COMB FROM   TO Type
0   PNR1  MAA  BLR    R
1   PNR1  BLR  MAA    R
2  PNR11  DEL  MAA    O
3   PNR2  TRV  HYD    R
4   PNR2  HYD  TRV    R