Pandas 一个索引级别的多索引匹配

Question

我有一个 pandas MultiIndex 对象，其中第一级是整数的常规递增索引，第二级包含其他整数，这些整数可能会或可能不会重复不同的 'frst' 索引值：

lst = list(filter(lambda x: x[1]%5 == x[0] or x[1]%4 == x[0],[(i,j) for i in range(5) for j in range(0, 20, 2)]))
mi = pd.MultiIndex.from_tuples(lst).rename(['frst', 'scnd'])
# mi = MultiIndex([(0,  0),(0,  4),(0,  8),(0, 10),(0, 12),(0, 16),(1,  6),(1, 16),(2,  2),(2,  6),(2, 10),(2, 12),(2, 14),(2, 18),(3,  8),(3, 18),(4,  4),(4, 14)], names=['frst', 'scnd'])

对于给定的 frst 值（例如 frst_idx = 0）和一些 shift，我需要找到 frst 为 frst_idx+shift 的所有索引，并且scnd 在 frst_idx 和 frst_idx+shift 之间共享。

例如：

frst_idx = 0, shift = 3 应该输出 [8] 因为上面的 MultiIndex 包含 (0, 8) 和 (3, 8).
frst_idx = 1, shift = 1 应该输出 [6] 因为 (1, 6) 和 (2, 6) 都在索引

所以我希望有一个函数可以接受这些参数和 return 所有匹配的 scnd 值的 pd.Series:

my_func(multi_index=mi, frst_idx=0, shift=3) ==> pd.Series([8])

迭代执行此操作非常昂贵 (O(n^2))，我希望有一些 pandas 魔法可以更快地执行此操作。

Answer 1

我找到了以下解决方案：

# reminder: $mi is a MultiIndex, mi.names = ['frst', 'scnd']
# assume some integer values for $frst_idx1, $shift

scnd_indices1 = mi[mi.get_level_values('frst') == frst_idx1].drop_level('frst')

frst_idx2 = frst_idx1 + shift
scnd_indices2 = mi[mi.get_level_values('frst') == frst_idx2].drop_level('frst')

result = scnd_indices1.intersection(scnd_indices2).to_series().reset_index(drop=True)

Pandas 一个索引级别的多索引匹配

Pandas MultiIndex match on one index level

python

pandas

multi-index