groupby 内的组合 pandas

Combinations within groupby pandas

我有以下数据框:

import pandas as pd

data = {'id': ['A', 'A', 'A', 'B', 'B', 'B', 'B'],
        'location':['Milan', 'Paris', 'New York', 'Rome', 'Los Angeles', 'Berlin', 'Madrid'],
        'year': [2003,2004,2005, 2003, 2004, 2004, 2005]}

data = pd.DataFrame(data)

对于每个 groupby('id'),我想要 year t 的城市 iyear t-1, t-2, ..., t-n 的所有城市之间的组合].

期望的输出:

data = {'id': ['A', 'A', 'A', 'A',
               'B', 'B', 'B', 'B', 'B', 'B'],
        'location':['Milan', 'Paris', 'New York', 'New York',
                    'Rome', 'Los Angeles', 'Berlin', 'Madrid','Madrid', 'Madrid'],
        'year': [2003, 2004, 2005, 2005,
                 2003, 2004, 2004, 2005, 2005, 2005],
       'comb': ['NaN', 'Milan', 'Milan','Paris',  
                'NaN', 'Rome', 'Rome', 'Rome','Los Angeles', 'Berlin']}

data = pd.DataFrame(data)

使用字典按年份获取位置

data = {'id': ['A', 'A', 'A', 'B', 'B', 'B', 'B'],
        'location':['Milan', 'Paris', 'New York', 'Rome', 'Los Angeles', 'Berlin', 'Madrid'],
        'year': [2003,2004,2005, 2003, 2004, 2004, 2005]}

df = pd.DataFrame(data)
print(df)

locations_by_year = {}
for year in df['year'].unique():
    locations_by_year[year] = df[df['year'] == year]['location'].unique()
    
print(locations_by_year)

输出:

{2003: array(['Milan', 'Rome'], dtype=object), 2004: array(['Paris', 'Los Angeles', 'Berlin'], dtype=object), 2005: array(['New York', 'Madrid'], dtype=object)}

年份位置:

df_grouped = df.groupby(['location'])
for name, group in df_grouped:
    print(name)
    print(group)

生成完整的笛卡尔积(原始数据帧所有行的所有组合)。然后按 df.year_comb < df.year 筛选。这也将删除每个 id 的第一年的行。如果需要,这些可以是 re-added 以在输出 df 中生成具有 NaN 值的行。

df = (pd.merge(data, data.rename(columns={"location": "comb", "year": "year_comb"}), on=["id"])
            .loc[lambda df: (df.year_comb < df.year)]
            .drop(["year_comb"], axis=1)
            )
# re-append the first years
data_first_years = data.sort_values(["year"]).groupby("id").first().reset_index()
df.append(data_first_years).sort_values(["id", "year"]).reset_index(drop=True)


# out:
  id     location  year         comb
0  A        Milan  2003          NaN
1  A        Paris  2004        Milan
2  A     New York  2005        Milan
3  A     New York  2005        Paris
4  B         Rome  2003          NaN
5  B  Los Angeles  2004         Rome
6  B       Berlin  2004         Rome
7  B       Madrid  2005         Rome
8  B       Madrid  2005  Los Angeles
9  B       Madrid  2005       Berlin

自合并,再查询:

N = 2
out = (data.merge(data, on='id', suffixes=['','_comb'])
           .query('0< year - year_comb <= @N')
      )

输出:

   id     location  year location_comb  year_comb
3   A        Paris  2004         Milan       2003
6   A     New York  2005         Milan       2003
7   A     New York  2005         Paris       2004
13  B  Los Angeles  2004          Rome       2003
17  B       Berlin  2004          Rome       2003
21  B       Madrid  2005          Rome       2003
22  B       Madrid  2005   Los Angeles       2004
23  B       Madrid  2005        Berlin       2004

注意:上面没有包含每个id的第一个位置,可以通过df.drop_duplicates('id')获取。所以你的最终输出将是

out = pd.concat([data.merge(data, on='id', suffixes=['','_comb'])
                     .query('0< year - year_comb <= @N'),
                 data.sort_values('year').drop_duplicates('id')] 
      )