Python：如何使用 Numpy 向量化我的拆分函数

Question

我在 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]))

Answer 1

这对于单个数组很容易：只需将您的序列设为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.array 中的tuple。该函数检查您的序列是否具有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 在 cmets 中指出的那样，这种转换比 Python list 或 tuple 本身的拆分花费的时间更长。想想你的数据在哪里，以及将它们放在 NumPy 数组中的什么位置。那些必须具有矩形形状。一旦你有了这个，你就可以编写一个合适的zigzag 版本。