All Courses

How to split numpy array with mask?

By Asif92, a year ago
  • Bookmark
0

I want to split array into array with mask and index

like below

a = array([ 0,  1,  2,  3,  4, 5]))  
b = [0,2,3]  

into

c = array([[0, 2, 3], [1, 3, 4], [2, 4, 5]])  

can I do this without loop?

Edit:

More examples...

Say, we have an array a with shape [10, 10, 10]

where a[x, y, :] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

Now given the mask b = [0, 3, 7]

I want the output to be an array c with shape [10, 10, 3, 3]

where c[x, y, :, :] = [[0, 3, 7], [1, 4, 8], [2, 5, 9]]

Numpy
Array
Slicing
1 Answer
0
Pritam Shinde

the answer is:

from numpy import array, arange
a = arange(2*2*10).reshape((2,2,10)) # some example input 
b = array([0,2,3])                   # the 'template' to extract
shifts = arange(3).reshape(-1,1)     # 3 is the number of repeats
indexer = b+shifts                   # broadcasted sum makes a matrix
c = a[:,:,indexer]                   # extract

This will take the b array as a kind of template, and repeat it with a certain shift. Finally, it will extract those entries from every array a[i,j,:] into c[i,j,:,:]. Otput from above is:

print(a)

[[[ 0  1  2  3  4  5  6  7  8  9]
  [10 11 12 13 14 15 16 17 18 19]]
 [[20 21 22 23 24 25 26 27 28 29]
  [30 31 32 33 34 35 36 37 38 39]]]

print(c)

[[[[ 0  2  3]
   [ 1  3  4]
   [ 2  4  5]]
  [[10 12 13]
   [11 13 14]
   [12 14 15]]]
 [[[20 22 23]
   [21 23 24]
   [22 24 25]]
  [[30 32 33]
   [31 33 34]
   [32 34 35]]]]


Your Answer

Webinars

Why You Should Learn Data Science in 2022?

Dec 15th (7:00 PM) 277 Registered
More webinars

Related Discussions

Running random forest algorithm with one variable

View More