Vectorize numpy数组for循环

时间:2012-12-18 17:21:14

标签: python numpy vectorization

我正在试图弄清楚如何对以下循环进行矢量化:

for i in range(1,size):
    if a[i] < a[i-1]:
        b[i] = a[i]
    else: b[i] = b[i-1]

b是与a相同大小的(大)数组。 我可以用

numpy.where(a[1:]<a[:-1])

替换if语句,但是如何同时替换else语句?

2 个答案:

答案 0 :(得分:2)

我想你想要这样的东西:

import numpy as np

def foo(a, b):
    # cond is a boolean array marking where the condition is met
    cond = a[1:] < a[:-1]
    cond = np.insert(cond, 0, False)
    # values is an array of the items in from a that will be used to fill b
    values = a[cond]
    values = np.insert(values, 0, b[0])
    # labels is an array of increasing indices into values
    label = cond.cumsum()
    b[:] = values[label]

答案 1 :(得分:1)

来自docs

  

numpy.where(condition[, x, y])

     
    

根据条件从x或y返回元素。

  

因此,您可以简化包含以下内容的循环:

if cond:
    a[i] = b[i]
else:
    a[i] = c[i]

a = numpy.where(cond, a, b)

但是,我不认为你的例子可以被矢量化,因为每个元素都依赖于前一个元素。