slice a 3d numpy array using a 2d numpy array

Question

Is it possible to slice a 3d array using a 2d array. Im assuming it can be done but would require that you have to specify the axis?

If I have 3 arrays, such that:

A = [[1,2,3,4,5],
     [1,3,5,7,9],
     [5,4,3,2,1]] # shape (3,5)

B1 = [[1],
      [2],
      [3]] # shape (3, 1) 

B2 = [[4],
      [3],
      [4]] # shape (3,1)

Is its possible to slice A using B1 an B2 like:

Out = A[B1:B2]

so that it would return me:

Out = [[2,3,4,5],
       [5, 7],
       [2, 1]]

or would this not work if the slices created arrays in Out of different lengths?

gsmafra · Accepted Answer · 2016-03-01 20:12:15Z

4

Numpy is optimized for homogeneous arrays of numbers with fixed dimensions, so it does not support varying row or column sizes.

However you can achieve what you want by using a list of arrays:

Out = [A[i, B1[i]:B2[i]+1] for i in range(len(B1))]

answered Mar 1, 2016 at 20:12

gsmafra

2,50422 silver badges27 bronze badges

Sign up to request clarification or add additional context in comments.

4 Comments

GeoMonkey Over a year ago

This is really neat but its not quite getting me what I want, it keeps returning the error that the slice is invalid. It did however get me thinking and whilst not as efficient I can get it to work within a loop using 'for i in range(len(B1))' followed by out.append([A[i][B1[i]:[B2[i]]]). Its probably not as efficient but it works for what I need. It doesn't appear to be too slow

gsmafra Over a year ago

I've tested the code myself and found no problems. Are you sure you tried this on numpy arrays of shape (3,5), (3,1) and (3,1)?

GeoMonkey Over a year ago

Admittedly, my arrays are not the same shape as this example case. A is (40000,432), B1 and B2 are (40000, 1). I assumed the example I provided would just be a scaled down version, apologies if not.

gsmafra Over a year ago

It should work anyway, as long as B2[i] is higher or equal than B1[i] for every i (and both should be integers of course)

Divakar · Accepted Answer · 2016-03-01 21:30:43Z

Here's one to vectorization -

n_range = np.arange(A.shape[1])
elems = A[(n_range >= B1) & (n_range <= B2)]      
idx = (B2 - B1 + 1).ravel().cumsum()
out = np.split(elems,idx)[:-1]

The trick is to use broadcasting to create a mask of elements to be selected for the output. Then, splitting the array of those elements at specified positions to get list of arrays.

Sample input, output -

In [37]: A
Out[37]: 
array([[1, 2, 3, 4, 5],
       [1, 3, 5, 7, 9],
       [5, 4, 3, 2, 1]])

In [38]: B1
Out[38]: 
array([[1],
       [2],
       [3]])

In [39]: B2
Out[39]: 
array([[4],
       [3],
       [4]])

In [40]: out
Out[40]: [array([2, 3, 4, 5]), array([5, 7]), array([2, 1])]
# Please note that the o/p is a list of arrays

hpaulj · Accepted Answer · 2016-03-01 21:50:38Z

Your desired result has a different number of terms in each row - that's a strong indicator that a fully vectorized solution is not possible. It is not doing the same thing for each row or each column.

Secondly, n:m translates to slice(n,m). slice only takes integers, not lists or arrays.

The obvious solution is some sort of iteration over rows:

In [474]: A = np.array([[1,2,3,4,5],
          [1,3,5,7,9],
          [5,4,3,2,1]]) # shape (3,5)

In [475]: B1=[1,2,3]  # no point in making these 2d

In [476]: B2=[5,4,5]  # corrected values

In [477]: [a[b1:b2] for a,b1,b2 in zip(A,B1,B2)]
Out[477]: [array([2, 3, 4, 5]), array([5, 7]), array([2, 1])]

This solution works just as well if A is a nested list

In [479]: [a[b1:b2] for a,b1,b2 in zip(A.tolist(),B1,B2)]
Out[479]: [[2, 3, 4, 5], [5, 7], [2, 1]]

The 2 lists could also be converted to an array of 1d indices, and then select values from A.ravel(). That would produce a 1d array, e.g.

array([2, 3, 4, 5, 5, 7, 2, 1]

which in theory could be np.split - but recent experience with other questions indicates that this doesn't save much time.

If the length of the row selections were all the same we can get a 2d array. Iterative version taking 2 elements per row:

In [482]: np.array([a[b1:b1+2] for a,b1 in zip(A,B1)])
Out[482]: 
array([[2, 3],
       [5, 7],
       [2, 1]])

I've discussed in earlier SO questions how produce this sort of result with one indexing operation.

On what slice accepts:

In [486]: slice([1,2],[3,4]).indices(10)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-486-0c3514e61cf6> in <module>()
----> 1 slice([1,2],[3,4]).indices(10)

TypeError: slice indices must be integers or None or have an __index__ method

'vectorized' ravel indexing

In [505]: B=np.array([B1,B2])    
In [506]: bb=A.shape[1]*np.arange(3)+B
In [508]: ri =np.r_[tuple([slice(i,j) for i,j in bb.T])]
# or np.concatenate([np.arange(i,j) for i,j in bb.T])

In [509]: ri
Out[509]: array([ 1,  2,  3,  4,  7,  8, 13, 14])

In [510]: A.ravel()[ri]
Out[510]: array([2, 3, 4, 5, 5, 7, 2, 1])

It still has an iteration - to generate the slices that go into np.r_ (which expands them into a single indexing array)

Collectives™ on Stack Overflow

slice a 3d numpy array using a 2d numpy array

3 Answers 3

4 Comments

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

3 Answers 3

4 Comments

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related