使用pandas groupby创建包含父数据帧的所有列的新数据框

Question

编辑：REWROTE整个问题

我根据经度和纬度编写了一个与商店位置列表匹配的程序与另一个医生位置列表。

每个商店位置都拥有唯一的ID。我想找到具有最小距离值的每个唯一ID（在这种情况下为单位）的3行。我希望这些行包含原始df中的所有信息。

使用此声明，我可以找到单位编号及其对应的距离，但没有别的。

df.groupby(['unit'])['distance'].nsmallest(3)

原始数据框示例

unit      store_name      distance     doctor_name
123       taqueria         0.3          dr. a
123       taqueria         0.35         dr. b
123       taqueria         0.2          dr. c
123       taqueria       0.5          dr. smith
123       taqueria       0.4          dr. q
123       taqueria       0.1          dr. who
456       carniceria         0.5          dr. a
456       carniceria         0.6         dr. b
456       carniceria         0.3          dr. c
456       carniceria       0.1         dr. smith
456       carniceria       0.2          dr. q
456       carniceria       0.05          dr. who

预期输出

unit      store_name      distance     doctor_name
123       taqueria         0.3          dr. a
123       taqueria         0.2          dr. c
123       taqueria       0.1          dr. who
456       carniceria       0.1         dr. smith
456       carniceria       0.2          dr. q
456       carniceria       0.05          dr. who

预期输出会将数据框缩减为按单位分组的三个最近 doctor_names

Answer 1

pandas

模糊和可爱，也很快

df.sort_values(['unit', 'distance']).groupby('unit').head(3)

    unit  store_name  distance doctor_name
5    123    taqueria      0.10     dr. who
2    123    taqueria      0.20       dr. c
0    123    taqueria      0.30       dr. a
11   456  carniceria      0.05     dr. who
9    456  carniceria      0.10   dr. smith
10   456  carniceria      0.20       dr. q

numpy

快速

v = df.values
u = df.unit.values
d = df.distance.values
s = d.size
r = np.arange(s)
a = np.lexsort([d, u])
u2 = u[a]

b = (np.where(np.append(True, u2[1:] != u2[:-1]))[0])
delta = np.diff(np.append(b, s))
mask = ((r - b.repeat(delta)) < 3)
pd.DataFrame(v[a][mask], df.index[a][mask], df.columns)

   unit  store_name distance doctor_name
5   123    taqueria      0.1     dr. who
2   123    taqueria      0.2       dr. c
0   123    taqueria      0.3       dr. a
11  456  carniceria     0.05     dr. who
9   456  carniceria      0.1   dr. smith
10  456  carniceria      0.2       dr. q

计时小数据

计时大数据

    

Answer 2

首先使用set_index（）

df = df.set_index(['store_name', 'doctor_name'])    
df1 = df.groupby(['unit'])['distance'].nsmallest(3).reset_index()

你得到了

    unit    store_name  doctor_name distance
0   123     taqueria    dr.who      0.10
1   123     taqueria    dr.c        0.20
2   123     taqueria    dr.a        0.30
3   456     carniceria  dr.who      0.05
4   456     carniceria  dr.smith    0.10
5   456     carniceria  dr.q        0.20