确定两个列表/数组的混洗索引

Question

作为挑战,我给自己这个问题:

给定2个列表A和B,其中B是A的混乱版本,其想法是找出混洗的索引.

例如:

A = [10, 40, 30, 2]
B = [30, 2, 10, 40]

result = [2,   3,    0,      1] 
        A[2]  A[3]   A[0]  A[1]
        ||     ||     ||    ||
        30      2     10    40

请注意,相同元素的关系可以任意解决.

我想出了一个解决方案,涉及使用字典来存储索引.这个问题有什么其他可能的解决方案？使用库的解决方案也有效.Numpy,pandas,一切都很好.

Answer 1

我们可以使用np.searchsorted它的可选sorter参数 -

sidx = np.argsort(B)
out = sidx[np.searchsorted(B,A, sorter=sidx)]

样品运行 -

In [19]: A = [10, 40, 30, 2, 40]
    ...: B = [30, 2, 10, 40]
    ...: 

In [20]: sidx = np.argsort(B)

In [21]: sidx[np.searchsorted(B,A, sorter=sidx)]
Out[21]: array([2, 3, 0, 1, 3])

Answer 2

作为对当前解决方案的改进，您可以使用collections.defaultdict和避免dict.setdefault：

from collections import defaultdict

A = [10, 40, 30, 2]
B = [30, 2, 10, 40]

idx = defaultdict(list)
for i, l in enumerate(A):
    idx[l].append(i)

res = [idx[l].pop() for l in B]
print(res)

以下是使用给定示例输入的两种方法的时序：

用于测试的脚本

from timeit import timeit


setup = """
from collections import defaultdict;
idx1 = defaultdict(list); idx2 = {}
A = [10, 40, 30, 2]
B = [30, 2, 10, 40]
"""

me = """
for i, l in enumerate(A):
    idx1[l].append(i)
res = [idx1[l].pop() for l in B]
"""

coldspeed = """
for i, l in enumerate(A):
    idx2.setdefault(l, []).append(i)
res = [idx2[l].pop() for l in B]
"""

print(timeit(setup=setup, stmt=me))
print(timeit(setup=setup, stmt=coldspeed))

结果

original: 2.601998388010543
modified: 2.0607256239745766

所以看来使用defaultdict实际上会产生轻微的速度提升。这实际上使得since虽然sincedefaultdict是用C而不是Python实现的。更不用说原始解决方案的属性查找idx.setdefault1成本高昂。