Python / NumPy第一次出现子数组


问题内容

在Python或NumPy中,找出子数组首次出现的最佳方法是什么?

例如,我有

a = [1, 2, 3, 4, 5, 6]
b = [2, 3, 4]

找出b在a中出现的最快方法(运行时)是什么?我了解字符串非常简单,但是对于列表或numpy ndarray呢?

非常感谢!

[编辑]我更喜欢numpy解决方案,因为根据我的经验,numpy向量化比Python列表理解要快得多。同时,大数组很大,因此我不想将其转换为字符串。那将(太)长。


问题答案:

我假设您正在寻找特定于numpy的解决方案,而不是简单的列表理解或for循环。一种方法可能是使用滚动窗口技术来搜索适当大小的窗口。这是rolling_window函数:

>>> def rolling_window(a, size):
...     shape = a.shape[:-1] + (a.shape[-1] - size + 1, size)
...     strides = a.strides + (a. strides[-1],)
...     return numpy.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
...

然后你可以做类似的事情

>>> a = numpy.arange(10)
>>> numpy.random.shuffle(a)
>>> a
array([7, 3, 6, 8, 4, 0, 9, 2, 1, 5])
>>> rolling_window(a, 3) == [8, 4, 0]
array([[False, False, False],
       [False, False, False],
       [False, False, False],
       [ True,  True,  True],
       [False, False, False],
       [False, False, False],
       [False, False, False],
       [False, False, False]], dtype=bool)

为了使它真正有用,您必须使用沿轴1减小它all

>>> numpy.all(rolling_window(a, 3) == [8, 4, 0], axis=1)
array([False, False, False,  True, False, False, False, False], dtype=bool)

然后您可以使用它,但是您将使用布尔数组。一种获取索引的简单方法:

>>> bool_indices = numpy.all(rolling_window(a, 3) == [8, 4, 0], axis=1)
>>> numpy.mgrid[0:len(bool_indices)][bool_indices]
array([3])

对于列表,您可以调整这些滚动窗口迭代器之一以使用类似的方法。

对于 非常 大的数组和子数组,可以这样保存内存:

>>> windows = rolling_window(a, 3)
>>> sub = [8, 4, 0]
>>> hits = numpy.ones((len(a) - len(sub) + 1,), dtype=bool)
>>> for i, x in enumerate(sub):
...     hits &= numpy.in1d(windows[:,i], [x])
... 
>>> hits
array([False, False, False,  True, False, False, False, False], dtype=bool)
>>> hits.nonzero()
(array([3]),)

On the other hand, this will probably be slower. How much slower isn’t clear
without testing; see Jamie‘s
answer for another memory-conserving option that has to check false positives.
I imagine that the speed difference between these two solutions will depend
heavily on the nature of the input.