在 pandas 中，根据表 B 中的条件获取表 A 中行索引的有效方法

Question

我正在尝试构建可用于时间序列建模的数据。 现在我有两张桌子：

表A:

Index UserID   SessionDate  
0      1       '2020-01-01'  
1      1       '2020-01-03'
2      2       '2020-03-01'
3      2       '2020-03-02'
4      3       '2020-01-05'

表B：

Index UserID   SnapshotDate  
0      1       '2020-01-01'  
1      1       '2020-01-02'
2      2       '2020-03-01'
3      2       '2020-03-02'
4      3       '2020-01-01'

因此，对于每个用户，在表B 中的每个快照日期，如果满足该用户的会话日期小于或等于快照日期，请给我表A 中的相应索引。

我试过使用 apply 函数

def index_search(x, df):
    user = x['UserID']
    snap_date = x['SnapshotDate']
    dd = df[df.UserID==user]
    ix = dd[dd.SessionDate <= snap_date].index.values
    return ix

idx = B.apply(index_search,df=A, axis=1)

但它很慢（我的数据集很大），所以我想知道有没有更有效的方法？

Answer 1

你可以试试numpy broadcasting：

x1, y1 = A.to_numpy().T
x2, y2 = B.to_numpy().T
mask = (x2[:, None] == x1) & (y2[:, None] >= y1)
idx = [A.index[m].tolist() for m in mask]

---

print(idx)

[[0], [0], [2], [2, 3], []]