Use numpy array to do conditional operations on another array-CodePudding

Let's say I have 2 arrays:

a = np.array([2, 2, 0, 0, 2, 1, 0, 0, 0, 0, 3, 0, 1, 0, 0, 2])

b = np.array([0, 0.5, 0.25, 0.9])

What I would like to do, is take the value in array b and multiple it to the values in array a, based on it's index.

So the first value in array a is 2. I want the value in array b at that index position to be multiplied by that value. So in array b, index postion 2's value is 0.25, so multiple that value (2) in array a by 0.25.

I know it can be done with iteration, but I'm trying to figure out how it's done elmentwise operations.

Here's the iteration way that I've done:

result = np.array([])
for idx in a:
    result = np.append(result, (b[idx] * idx))

To get the result:

print(result)
[0.5 0.5 0.  0.  0.5 0.5 0.  0.  0.  0.  2.7 0.  0.5 0.  0.  0.5]

What's an elementwise equivalent?

CodePudding user response：

Integer arrays can be used as indices in numpy. As a consequence, you can simply do something like this

b[a] * a

EDIT:

Just for completeness, your iterative solution triggers a new memory allocation every time append is called (see the 'returns' section of this page). Since you already now the shape of your output (i.e. a.shape), it's much better to allocate the output array in advance, e.g. result = np.empty(a.shape) and then go through the cycle.

CodePudding user response：

So there are a few ways to do this, but if you want purely element-wise operations you could do the following:

Before getting the result, each element of b is transformed by its index. So create another vector n.

n = np.arange(len(b)) * b

# In the example, n now equals [0. , 0.5, 0.5, 2.7]

# then the result is just n indexed by a

result = n[a]

# result = [0.5, 0.5, 0. , 0. , 0.5, 0.5, 0. , 0. , 0. , 0. , 2.7, 0. , 0.5, 0. , 0. , 0.5]