根据另一列中的值来限制整个数据帧行中的公用值

时间:2019-02-28 22:40:06

标签: python pandas dictionary combinations itertools

以下面的示例数据框为例,如何找到相同兴趣的相同位置的用户?数据是-

userid   interest    location
 1       [A, B]          Z
 2       [A, C, B]       Y
 3       [B, D]          Z
 4       [A, C]          Y
 5       [A, B, D]       Z

输出应为-

userid    relativeid  common interest  location
  1          3           [B]              Z
  1          5           [A, B]           Z
  2          4           [A,C]            Y

到目前为止,我为每个位置都创建了单独的数据框,如下所示-

userid   interest    location
 1       [A, B]          Z
 3       [B, D]          Z
 5       [A, B, D]       Z

代码-

dictionary = df.set_index('userid')['interest'].map(set).to_dict()
dictionary

out = pd.DataFrame(list(itertools.combinations(df.userid, 2)), columns=['userid', 'relative_id']) 
out['common_interest'] = [list(dictionary[x] & dictionary[y]) for x, y in out.values]
out

但这给我的输出没有location列。

userid    relativeid  common interest 
  1          3           [B]              
  1          5           [A, B]

问题: 1)如何修改此代码以获取输出中的location列? 2)有没有一种方法可以不根据位置将原始数据帧分为多个数据帧?

2 个答案:

答案 0 :(得分:0)

这是一种可能的解决方案。我注意到添加的代码。仅创建了一个额外的字典来保留每个用户ID的位置信息,该函数会删除用户的组合(如果他们不共享相同的位置),最后一行使用该相同的位置字典在页面上创建位置列最终数据框。

import itertools

users_df = pd.DataFrame({'userid':[1,2,3,4,5],'interest':[['A','B'],['A','C','B'],['B','D'],['A','C'],['A','B','D']],
                        'location':['Z','Y','Z','Y','Z']})

#new code: location dictionary
loc_dict = users_df.set_index('userid')['location'].to_dict()

#new code: function that removes userid combinations when locations are different
def restrict_users(all_combs):
    return [comb for comb in all_combs if loc_dict[comb[0]] == loc_dict[comb[1]]]

dictionary = users_df.set_index('userid')['interest'].map(set).to_dict()

#new function applied below
out = pd.DataFrame(restrict_users(list(itertools.combinations(users_df.userid, 2))), columns=['userid', 'relative_id']) 
out['common_interest'] = [list(dictionary[x] & dictionary[y]) for x, y in out.values]
#location column added to the dataframe
out['location'] = out['userid'].map(loc_dict)

out

Output

答案 1 :(得分:0)

这是我的解决方案,而无需创建子数据帧。虽然看起来有点沉重。感谢将@AlexK作为框架的初始内容。

import pandas as pd
import itertools

df = pd.DataFrame({'userid':[1,2,3,4,5],'interest':[['A','B'],['A','C','B'],['B','D'],['A','C'],['A','B','D']],
                        'location':['Z','Y','Z','Y','Z']})
# Builds a dictionary of location as key and a list of index of users in df as value
idxlocation = df.groupby('location').apply(lambda x: x.index.values).to_dict()

new_frame = []
for k, v  in idxlocation.items():
    for i in itertools.combinations(v, 2):
        userid = df.loc[i[0], 'userid']
        relativeid = df.loc[i[1], 'userid']
        new_frame.append((userid, relativeid, [j for j in set(df.loc[i[0], 'interest']).intersection(set(df.loc[i[1], 'interest']))], k))

out = pd.DataFrame(new_frame)
out.columns = ['userid', 'relative_id', 'common_interest', 'location']

>>>out
   userid  relative_id common_interest location
0       2            4          [A, C]        Y
1       1            3             [B]        Z
2       1            5          [A, B]        Z
3       3            5          [D, B]        Z