Pandas：根据相似的列和（列）列表元素分配组 ID

Question

我的数据包含人名和他们居住的城市列表。我想按照以下条件将它们组合在一起：

• 1. first_name 和 last_name 相同
• 1. 或（如果 1. 不成立）他们的last_name 是相同的，并且他们至少住在一个相同的city 中。

结果应该是一个新列，指示每个人所属的组 ID。
DataFrame df 如下所示：

>>> df
      person_id  last_name first_name                     cities
    0       112     Dorsey      Nancy         [Moscow, New York]
    1       113     Harper        Max  [Munich, Paris, Shanghai]
    2       114    Mueller        Max    [New York, Los Angeles]
    3       115     Dorsey      Nancy          [New York, Miami]
    4       116     Harper    Maxwell            [Munich, Miami]

新的数据框df_id 应该如下所示。 id 的顺序无关紧要（即，哪个组得到 id=1），但只有满足条件 1 或 2 的观察才能得到相同的 id。

>>> df_id
      person_id  last_name first_name                     cities  id
    0       112     Dorsey      Nancy         [Moscow, New York]   1
    1       113     Harper        Max  [Munich, Paris, Shanghai]   2
    2       114    Mueller        Max    [New York, Los Angeles]   3
    3       115     Dorsey      Nancy          [New York, Miami]   1
    4       116     Harper    Maxwell            [Munich, Miami]   2

我当前的代码：

df= df.reset_index(drop=True)
#explode lists to rows
df_exploded = df.explode('cities')
# define id_counter and dictionary to person_id to id
id_counter = 1
id_matched = dict()

# define id function
def match_id(df):
  global id_counter

  # check if person_id already matched
  if df['person_id'] not in id_matched.keys():
  # get all persons with similar names (condition 1)
    select = df_expanded[(df_expanded['first_name']==df['first_name']) & df_expanded['last_name']==df['last_name'])]
    # get all persons with same last_name and city (condition 2)
    if select.empty:
      select_2 = df_expanded[(df_expanded['last_name']==df['last_name']) & (df_expanded['cities'] in df['cities'])]
      # create new id for this specific person
      if select_2.empty:
        id_matched[df['person_id']] = id_counter
      # create new id for group of person and record in dictionary
      else:
        select_list = select_2.unique().tolist()
        select_list.append(df['person_id'])
        for i in select_list:
          id_matched[i] = id_counter 
    # create new id for group of person and record in dictionary
    else:
      select_list = select.unique().tolist()
      select_list.append(df['person_id'])
      for i in select_list:
        id_matched[i] = id_counter
    # set next id
    id_counter += 1

# run function
df = df.progress_apply(match_id, axis=1)

# convert dict to DataFrame
df_id_matched = pd.DataFrame.from_dict(id_matched, orient='index', columns['id'])
                            .rename_axis('person_id').reset_index()

# merge back together with df to create df_id

有没有人有更有效的方法来执行这项任务？数据集很大，需要几天时间...

提前致谢！

Answer 1

• cond1 设置 id 具有相同的 first_name + last_name

df['name'] = df['first_name'].fillna('') + '_' + df['last_name'].fillna('')
obj_name = df['name'].unique()
name_map = dict(zip(obj_name, np.arange(1, len(obj_name)+1)))
df['tag'] = df['name'].map(name_map)

• cond2，找同城同姓

obj_last_name_city = df[['last_name', 'cities']].explode('cities')
cond = obj_last_name_city.duplicated(keep=False)
objn = obj_last_name_city[cond].reset_index().groupby(['last_name','cities'])['index'].agg(list)

# sort_values is import if there one last_name with n same cities or reverse.
for i in objn.sort_values().tolist():
     df.loc[i[1:], 'tag'] = df.loc[i[0], 'tag']

# print(objn)

#     last_name  cities  
#     Dorsey     New York    [0, 3]
#     Harper     Munich      [1, 4]
#     Name: index, dtype: object

输出：

   person_id last_name first_name                     cities  tag
0        112    Dorsey      Nancy         [Moscow, New York]    1
1        113    Harper        Max  [Munich, Paris, Shanghai]    2
2        114   Mueller        Max    [New York, Los Angeles]    3
3        115    Dorsey      Nancy          [New York, Miami]    1
4        116    Harper    Maxwell            [Munich, Miami]    2

Answer 2

• 另一种解决方案
• 参考：merge-lists-that-share-common-elements

import networkx 
from networkx.algorithms.components.connected import connected_components
def to_graph(l):
    G = networkx.Graph()
    for part in l:
        # each sublist is a bunch of nodes
        G.add_nodes_from(part)
        # it also imlies a number of edges:
        G.add_edges_from(to_edges(part))
    return G

def to_edges(l):
    """ 
        treat `l` as a Graph and returns it's edges 
        to_edges(['a','b','c','d']) -> [(a,b), (b,c),(c,d)]
    """
    it = iter(l)
    last = next(it)

    for current in it:
        yield last, current
        last = current    

• 开始：

df = pd.DataFrame(
    [{'person_id': 112, 'last_name': 'Dorsey', 'first_name': 'Nancy', 'cities': ['Moscow', 'New York'], 'id': 1},
     {'person_id': 113, 'last_name': 'Harper', 'first_name': 'Max', 'cities': ['Munich', 'Paris', 'Shanghai'], 'id': 2},
     {'person_id': 114, 'last_name': 'Mueller', 'first_name': 'Max', 'cities': ['New York', 'Los Angeles'], 'id': 3}, 
     {'person_id': 115, 'last_name': 'Dorsey', 'first_name': 'Nancy', 'cities': ['New York', 'Miami'], 'id': 1},
     {'person_id': 116, 'last_name': 'Harper', 'first_name': 'Maxwell', 'cities': ['Munich', 'Miami'], 'id': 2}, 
     {'person_id': 117, 'last_name': 'Harper', 'first_name': 'Jan', 'cities': ['Miami'], 'id': 2}
    ]
)
df = df.reset_index()

# cond1
obj_name = df['first_name'].fillna('') + '_' + df['last_name'].fillna('')
obj1 = df.groupby(obj_name)['index'].agg(list).to_list()

# cond2
obj_last_name_city = df[['last_name', 'cities']].explode('cities')
cond = obj_last_name_city.duplicated(keep=False)
obj2 = obj_last_name_city[cond].reset_index().groupby(['last_name','cities'])['index'].agg(list).to_list()

L =(obj1 + obj2)
# [[1], [2], [4], [0, 3], [0, 3], [1, 4]]

# use networkx
G = to_graph(L)
obj = pd.Series(connected_components(G))
obj = obj.map(list).sort_values()

idx_map = (obj.reset_index(drop=True)
           .explode()
           .reset_index()
           .set_index(0)['index']
           .to_dict())

df['_id'] = df['index'].map(idx_map) + 1
df