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