np.where for 2d array, manipulate whole rows

--------------------------------------------------
Rise to the top 3% as a developer or hire one of them at Toptal: https://topt.al/25cXVn
--------------------------------------------------

Music by Eric Matyas
https://www.soundimage.org
Track title: Puzzle Meditation

--

Chapters
00:00 Np.Where For 2d Array, Manipulate Whole Rows
00:34 Answer 1 Score 2
00:45 Accepted Answer Score 2
01:06 Thank you

--

Full question
https://stackoverflow.com/questions/7158...

--

Content licensed under CC BY-SA
https://meta.stackexchange.com/help/lice...

--

Tags
#python #arrays #numpy #arraybroadcasting

#avk47

ANSWER 1

Score 2

IIUC you want something like this:

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

ACCEPTED ANSWER

Score 2

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