获取与同一列下的两个值匹配的记录(查找共享关联)



如何查询熊猫数据帧以获取与一列下不同值关联的记录?它本质上是一个共同查询:

data = {"id": ["a", "a", "a", "b", "b", "c", "c", "d", "e", "f", "f", "f"],
"x": [1, 2, 3, 1, 3, 5, 1, 7, 2, 4, 9, 11],
"y": [1985, 1986, 1987, 1985, 1987, 1990, 1985, 1994, 1985, 1989, 1993, 1993]}
df = pd.DataFrame(data)
print(df)
id   x     y
0   a   1  1985
1   a   2  1986
2   a   3  1987
3   b   1  1985
4   b   3  1987
5   c   5  1990
6   c   1  1985
7   d   7  1994
8   e   2  1985
9   f   4  1989
10  f   9  1993
11  f  11  1993

鉴于上述数据,我想编写一个函数f(id1, id2, year=None)来查找与 id1 和 id2 关联的x。因此,f('a', 'b')将返回对应于 {x: [1, 3], y: [1985, 1987]} 的数据帧

x     y
0  1  1985
1  3  1987 

因为 x=1 和 x=3 都与 'a' 和 'b' 相关联。

理想情况下,我希望能够以允许按"y"值过滤的方式执行此操作。此外,数据帧非常大,这是一个经常运行的实用程序函数,因此效率很重要。我想过使用不同的ID查询两次并合并:

In [15]: pd.merge(df.query('id=="a"'), df.query('id=="b"'), on='x')
Out[15]:
id_x  x   y_x id_y   y_y
0    a  1  1985    b  1985
1    a  3  1987    b  1987

但我有一种感觉,这不是最优雅和最有效的方式。

我修改了一点 moys 答案并添加了性能。 你可以做:

data = pd.DataFrame({"id": ["a", "a", "a", "b", "b", "c", "c", "d", "e", "f", "f", "f"],
"x": [1, 2, 3, 1, 3, 5, 1, 7, 2, 4, 9, 11],
"y": [1985, 1986, 1987, 1985, 1987, 1990, 1985, 1994, 1985, 1989, 1993, 1993]})
def finder(ida,idb,year=None):
g = data.groupby('id')
g1 = set(g["x"].get_group(ida))
g2 = set(g["x"].get_group(idb))
if year:
return data.loc[(data.x.isin(g1&g2)) & (data.id.isin([ida,idb])) & (data.y==year)]
else :
return data.loc[(data.x.isin(g1&g2)) & (data.id.isin([ida,idb])

它画

每个环路 4.42 ms ± 36.9 μs(平均 ± 标准偏差 7 次运行,每次 100 次循环(

相比之下,您当前的代码绘制

每个环路 10.2 ms ± 123 μs(7 次运行的平均标准±偏差,每次 100 次循环(

更新后的 moys 代码绘制:

每个环路 10.6 ms ± 60.3 μs(7 次运行的平均标准±差,每次 100 次循环(

IIUC,这应该可以满足您的主要需求。

def f(x,y):
g = df.groupby('id')
g1 = g["x"].get_group(x).to_list()
g2 = g["x"].get_group(y).to_list()
return [val for val in g1 if val in g2]
f('a', 'b')

输出

[1,3]


如果需要输出是数据帧,下面的代码应该可以工作。

def f(x,y):
g = df.groupby('id')
if g.get_group(x).shape[0] >= g.get_group(y).shape[0]:
res = g.get_group(x)[['x','y']].merge(g.get_group(y)[['x','y']], on=['x','y'],how = 'inner')
else:
res = g.get_group(y)[['x','y']].merge(g.get_group(x)[['x','y']], on=['x','y'],how = 'inner')
return res
f('a', 'b')

输出

x   y
0   1   1985
1   3   1987

最新更新