有效地在 Python 列表中查找索引（与 MATLAB 相比）

Question

我很难找到在 Python 列表中查找索引的有效解决方案。到目前为止，我测试过的所有解决方案都比 MATLAB 中的“查找”函数慢。我才刚刚开始使用 Python（因此，我不是很有经验）。

在 MATLAB 中，我会使用以下代码：

a = linspace(0, 1000, 1000); % monotonically increasing vector
b = 1000 * rand(1, 100); % 100 points I want to find in a
for i = 1 : numel(b)
    indices(i) = find(b(i) <= a, 1); % find the first index where b(i) <= a
end

如果我使用 MATLAB 的 arrayfun()，我可以稍微加快这个过程。 在 Python 中，我尝试了几种可能性。我用过

for i in xrange(0, len(b)):
   tmp = numpy.where(b[i] <= a)
   indices.append(tmp[0][0])

这需要很多时间，尤其是当 a 很大时。 如果 b 已排序，我可以使用

for i in xrange(0, len(b)):
    if(b[curr_idx] <= a[i]):
        indices.append(i)
        curr_idx += 1
    if(curr_idx >= len(b)):
        return indices
        break

这比 numpy.where() 解决方案快得多，因为我只需在列表中搜索一次，但这仍然比 MATLAB 解决方案慢。

谁能提出更好/更有效的解决方案？ 提前致谢。

Answer 1

试试numpy.searchsorted:

>> a = np.array([0, 1, 2, 3, 4, 5, 6, 7])
>> b = np.array([1, 2, 4, 3, 1, 0, 2, 9])
% sorting b "into" a
>> np.searchsorted(a, b, side='right')-1
array([1, 2, 4, 3, 1, 0, 2, 9])

您可能需要对 b 中超出 a 范围的值进行一些特殊处理，例如上例中的 9。 尽管如此，这应该比任何基于循环的方法都要快。

顺便说一句： 同样，MATLAB 中的histc 会比循环快得多。

编辑：

如果您想要获取b 最接近a 的索引，您应该可以使用相同的代码，只需修改a：

>> a_mod = 0.5*(a[:-1] + a[1:]) % take the centers between the elements in a
>> np.searchsorted(a_mod, np.array([0.9, 2.1, 4.2, 2.9, 1.1]), side='right')
array([1, 2, 4, 3, 1])

请注意，您可以删除-1，因为a_mod 比a 少一个元素。

Answer 2

使用 numpy 仅用于数字生成（不适用于矢量化）：

import numpy as np
a = np.linspace(0, 1000, 1000)
b = 1000 * np.random.rand(100)
indices = [next(i for i, ai in enumerate(a) if bi <= ai) for bi in b]

如果a.max() >= b.max() 如示例中所示，这将有效，否则将引发StopIteration，并且它仍然很慢（尽管这并不能像b(i) <= a 那样进行所有可能的比较）。

如果您需要将索引作为数组而不是列表，请在此之后使用np.array(indices)。如果需要优化，可以对b进行排序，只保留一个enumerate(a)，偷看而不是取最后一个元素。

你也可以尝试在 pypy 上不使用 numpy：

def igen(a, b):
    iterb = iter(b)
    bi = next(iterb)
    for i, ai in enumerate(a):
        while bi <= ai:
            yield i
            bi = next(iterb)
    i += 1 # Last bi are bigger than all ai
    yield i
    for unused in iterb:
        yield i

from random import random
a = (i * 1000. / 999. for i in xrange(43032500))
b = sorted(random() * 1000 for unused in xrange(3848))
indices = list(igen(a, b))

这个是基于使用该思想的生成器，并且 b 应该被排序。这将在所有ai 的bi > ai 时返回len(a)。

为了测试，我正在使用：

setup = """
from random import random

def igen(a, b):
    iterb = iter(b)
    bi = next(iterb)
    for i, ai in enumerate(a):
        while bi <= ai:
            yield i
            bi = next(iterb)
    i += 1 # Last bi are bigger than all ai
    yield i
    for unused in iterb:
        yield i
"""

program = """
a = (i * 1000. / 999. for i in xrange(43032500))
b = sorted(random() * 1000 for unused in xrange(3848))
indices = list(igen(a, b))
"""

# Python 2 and 3 compatibility
import sys
if sys.version_info.major == 3:
    program = program.replace("xrange", "range")

# Time it! =)
from timeit import timeit
print(timeit(program, setup, number=5000))

这意味着我在每个环境中都运行了 5000 次该算法。结果时间是所有试验 (program) 持续时间的总和（不是平均值）：

• 在 CPython 3.4.0 上，结果为 11.491293527011294（秒）
• 在 CPython 2.7.6 上，结果为 9.39319992065（秒）
• 在 Pypy 2.2.1 上，结果为 3.31203603745（秒）

更具体的版本信息：

• Linux 上的 Python 3.4.0（默认，2014 年 4 月 11 日，13:05:11）[GCC 4.8.2]
• Linux2 上的 Python 2.7.6（默认，2014 年 3 月 22 日，22:59:56）[GCC 4.8.2]
• Linux2 上的 Python 2.7.3（2.2.1+dfsg-1，2013 年 11 月 28 日，05:13:10）[PyPy 2.2.1 和 GCC 4.8.2]

现在与改编的“两个ifs”版本相同（代码如下）得到了结果：

• 在 CPython 3.4.0 上，结果为 13.03860338096274（秒）
• 在 CPython 2.7.6 上，结果为 10.7371659279（秒）
• 在 Pypy 2.2.1 上，结果为 2.88891601562（秒）

Pypy 找到了一种方法来优化您的版本，但仍然有一个区别，我已经测试了这个计算“a”的一次，而我的版本计算了 5000 次“a”。我运行的代码是：

setup = """
from random import random
a = [i * 1000. / 999. for i in xrange(43032500)]
"""

program = """
b = sorted(random() * 1000 for unused in xrange(3848))
curr_idx = 0
indices = []
for i in xrange(len(a)): # Why not for i, ai in enumerate(a)?
    if b[curr_idx] <= a[i]:
        indices.append(i)
        curr_idx += 1
    if curr_idx >= len(b):
        break
"""

# Python 2 and 3 compatibility
import sys
if sys.version_info.major == 3:
    setup = setup.replace("xrange", "range")
    program = program.replace("xrange", "range")

# Time it! =)
from timeit import timeit
print(timeit(program, setup, number=5000))

另一个版本只是将a 分配给program，而不是将其保留在setup，这样做Pypy 时间转到2102.06863689（是的，超过35 分钟）。将东西存储在一个巨大的列表上真的很慢。将程序开头更改为：

a = (i * 1000. / 999. for i in xrange(43032500)) # A generator expression
[...]
for i, ai in enumerate(a):
    if b[curr_idx] <= ai:
    [...]

使用 Pypy 让我们回到 3.11599397659 秒。在这个版本中，a 被创建了 5000 次，但从未存储在列表中。另一方面，函数外部“硬编码”的igen 版本在3.17516112328 秒上工作，其中setup 刚刚导入random 和program：

a = (i * 1000. / 999. for i in xrange(43032500))
b = sorted(random() * 1000 for unused in xrange(3848))
indices = []
iterb = iter(b)
try:
    bi = next(iterb)
    for i, ai in enumerate(a):
        while bi <= ai:
            indices.append(i)
            bi = next(iterb)
except StopIteration:
    pass
else:
    i += 1 # Last bi are bigger than all ai
    indices.append(i)
    for unused in iterb:
        indices.append(i)

无论如何，让A = len(a) 和B = len(b)，所以这些是O[A + B.log(B)] 算法（包括带有np.searchsorted 的@sebastian 解决方案）。另一方面，为所有对 (bi, ai) 计算 bi <= ai 是 O[b * a]，除非它进行一些内部优化以避免完全比较，否则 Matlab 解决方案应该渐近缓慢（但我没有 Matlab验证 =/)。作为比较的需要，我在 GNU Octave 上做了这个：

start = time;
a = linspace(0, 1000, 43032500);
b = 1000 * rand(1, 3848);
for i = 1 : numel(b)
    indices(i) = find(b(i) <= a, 1);
end
stop = time;

stop - start

这是 Python 使用此问题的原始代码执行 5000 次的过程，它发生在 203.16 秒（超过 3 分钟）内。

哦，但你在作弊！把那个“开始=时间；”在分配给“a”之后！

好的，没有人这么说，但我刚刚尝试过这样的改变。由于每个b(i) <= a 都是一个大小为 43032500 的向量，因此变化不大：202.83 秒。

还有 Numpy？！

Numpy 也必须存储数据。大多数情况下，它不适用于生成器（hstack 和 vstack 是例外）。但如果没有经验证据，我们无法确定哪个更快。让我们用 Numpy 1.8.1 运行它：

setup = """
import numpy as np
a = np.linspace(0., 1000., 43032500) # Don't count this time
"""

program = """
b = 1000 * np.random.rand(3848)
indices = np.searchsorted(a, b, side='right') - 1 # From @sebastian solution
indices[b > a[-1]] = len(a) # Big value correction (my improvement)
"""

# Time it! =)
from timeit import timeit
print(timeit(program, setup, number=5000))

• 在 CPython 2.7 上，9.81494688988 秒
• 在 CPython 3.4 上，9.831143222982064 秒

就是这样。 =)