CODEWITHSUNDEEP
pythonarraysfunctionnumpy
Ashok
Ashok

Reputation: 214

Python- Numpy array in function with conditions

I have a numpy array (say xs), for which I am writing a function to create another array (say ys), which has the same values as of xs until the first half of xs and two times of xs in the remaining half. for example, if xs=[0,1,2,3,4,5], the required output is [0,1,2,6,8,10]

I wrote the following function:

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

Which produces output: `array([0., 1., 2., 3., 4., 5.]), ie the second condition is not executed. Does anybody know how to fix it? Thanks in advance.

Upvotes: 1

Views: 530

Answers (2)

hyloop
hyloop

Reputation: 359

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]

Sorry, I did not find your bug, but I felt it can be implemented differently.

Upvotes: 2

jpp
jpp

Reputation: 164683

You can use vectorised operations instead of Python-level iteration. With the below method, we first copy the array and then multiply the second half of the array by 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]

Upvotes: 2

Related Questions