np.where for 2d array, manipulate whole rows

Question

I want to rebuild the following logic with numpy broadcasting function such as np.where: From a 2d array check per row if the first element satisfies a condition. If the condition is true then return the first three elements as a row, else the last three elements.

A short MWE in form of a for-loop which I want to circumvent:

import numpy as np
array = np.array([
    [1, 2, 3, 4],
    [1, 2, 4, 2],
    [2, 3, 4, 6]
])

new_array = np.zeros((array.shape[0], array.shape[1]-1))
for i, row in enumerate(array):
    if row[0] == 1: new_array[i] = row[:3]
    else: new_array[i] = row[-3:]

`np.where` returns `row, col` if nothing given as second parameter — Flash Thunder, Mar 23 '22 at 15:17

score 2 · Answer 1 · answered Mar 23 '22 at 15:15

2

IIUC you want something like this:

condition = array[:,0]==1
new_array[condition,:] = array[condition,:3]
new_array[~condition,:] = array[~condition,-3:]

answered Mar 23 '22 at 15:15

warped

8,947
3
22
49

in the third line you are missing two "~" in the conditon, no? Negating the condition – malwin Mar 23 '22 at 15:26
@malwin There are two `~` there. – warped Mar 23 '22 at 15:30

score 2 · Accepted Answer · answered Mar 23 '22 at 15:20

If you want to use np.where:

import numpy as np
array = np.array([
    [1, 2, 3, 4],
    [1, 2, 4, 2],
    [2, 3, 4, 6]
])

cond = array[:, 0] == 1
np.where(cond[:, None], array[:,:3], array[:,-3:])

output:

array([[1, 2, 3],
       [1, 2, 4],
       [3, 4, 6]])

EDIT

slightly more concise version:

np.where(array[:, [0]] == 1, array[:,:3], array[:,-3:])

np.where for 2d array, manipulate whole rows

2 Answers2