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))

第二个序列只是一个tuple of numpy.array。该函数检查您的序列是否具有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或list本身的拆分要长得多。考虑一下您的数据在哪里，在NumPy数组中有哪些数据是有意义的。必须是长方形的。一旦您有了它，您就可以编写一个合适的zigzag版本。