Python : 如何用 Numpy 向量化我的分割函数

Python : How to vectorize my split function with Numpy

我在 Stack 上找到了一个可以解决我问题的函数,但现在我想加快我的代码速度,因为我有很多列表要拆分。

我听说向量化函数是一个解决方案,所以尝试用 numpy 向量化我的函数,但它不起作用。

你能帮帮我吗?

原函数:

seq = ([1,1,5,1,5,5,1,5,1,1])

def zigzag(seq):
  return seq[::2], seq[1::2]

结果:

([1, 5, 5, 1, 1], [1, 1, 5, 5, 1])

我的矢量化尝试:

import numpy as np
seq = ([1, 1, 5, 1, 5, 5, 1, 5, 1, 1], [2, 2, 2, 3, 3, 3, 3, 2, 2, 2], [6, 3, 9, 2, 9, 4, 6, 3])

def zigzag(seq):
  return seq[::2], seq[1::2]

vecto = np.vectorize(zigzag)

vecto(seq)

想要的结果:

(([1, 5, 5, 1, 1], [1, 1, 5, 5, 1]), ([2, 2, 3, 3, 2], [2, 3, 3, 2, 2]), ([6, 9, 9, 6], [3, 2, 4, 3]))

这对于单个数组来说很容易:只需将您的序列设为 numpy.array,您的 zigzag 函数将在后台调用 C 代码。

def zigzag(seq):
  return seq[::2], seq[1::2]

seq = np.array([1,1,5,1,5,5,1,5,1,1])

result = zigzag(seq)
print(result)

结果:

(array([1, 5, 5, 1, 1]), array([1, 1, 5, 5, 1]))

对于多维情况,您会遇到列表长度不一样的问题。因此,您无法从中得到很好的 numpy.array。我建议你这样改编:

import numpy as np

def zigzag(seq):
    try:
        if len(seq.shape) == 1:
            return seq[::2], seq[1::2]
    except AttributeError:
        return [zigzag(x) for x in seq]

def main():
    options = _parse_args()

    seq = np.array([1,1,5,1,5,5,1,5,1,1])
    seq2 = (
        np.array([1, 1, 5, 1, 5, 5, 1, 5, 1, 1]),
        np.array([2, 2, 2, 3, 3, 3, 3, 2, 2, 2]),
        np.array([6, 3, 9, 2, 9, 4, 6, 3]),
    )

    print(zigzag(seq))
    print()
    print(zigzag(seq2))

第二个序列只是 numpy.arraytuple。该函数检查您的序列是否具有 shape 属性,这是一个很好的指标,表明它是 numpy.array。如果是这样,它会使用 NumPy 切片。如果它是一个元组,它只会为每个元素调用 zigzag 函数。

它为您的示例生成所需的输出:

(array([1, 5, 5, 1, 1]), array([1, 1, 5, 5, 1]))

[(array([1, 5, 5, 1, 1]), array([1, 1, 5, 5, 1])), (array([2, 2, 3, 3, 2]), array([2, 3, 3, 2, 2])), (array([6, 9, 9, 6]), array([3, 2, 4, 3]))]

然而,这不是一个完善的解决方案。您不想一直将 Python 列表和元组转换为 NumPy 数组。正如@hpaulj 在评论中指出的那样,这种转换比拆分 Python listtuple 本身花费的时间更长。想想你的数据在哪里,以及将它们放在 NumPy 数组中的什么地方是有意义的。那些必须具有矩形形状。一旦你有了这个,你就可以写一个合适的 zigzag 版本。