函数中带条件的Python Numpy数组

Question

我有一个numpy数组（比如xs），我正在为它编写一个函数来创建另一个数组（比如ys），这个数组的值与xs的值相同，直到xs的前半部分，其余部分是xs的两倍。例如，如果xs=[0,1,2,3,4,5]，则所需输出为[0,1,2,6,8,10]

我编写了以下函数：

import numpy as np
xs=np.arange(0,6,1)

def step(xs):
    ys1=np.array([]);ys2=np.array([])
    if xs.all() <=2:
        ys1=xs
    else:
        ys2=xs*2
    return np.concatenate((ys1,ys2))

print(xs,step(xs))

它产生输出：`array（[0,1,2,3,4,5]），即第二个条件不被执行。有人知道怎么修吗？提前谢谢。你知道吗

Answer 1

import numpy as np

xs=np.arange(0,6,1)

def f(a):
    it = np.nditer([a, None])
    for x, y in it:
        y[...] = x if x <= 2 else x * 2
    return it.operands[1]

print(f(xs))

---

[ 0  1  2  6  8 10]

抱歉，我没有找到你的错误，但我觉得它可以实现不同的。你知道吗

Answer 2

您可以使用向量化操作，而不是Python级迭代。使用下面的方法，我们首先复制数组，然后将数组的后半部分乘以2。你知道吗

import numpy as np

xs = np.arange(0,6,1)

def step(xs):
    arr = xs.copy()
    arr[int(len(arr)/2):] *= 2
    return arr

print(xs, step(xs))

[0 1 2 3 4 5] [ 0  1  2  6  8 10]