CODEWITHSUNDEEP
pythonarraysnumpyslice
GeoMonkey
GeoMonkey

Reputation: 1671

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

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?

Upvotes: 1

Views: 1190

Answers (3)

hpaulj
hpaulj

Reputation: 231665

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)

Upvotes: 0

Divakar
Divakar

Reputation: 221684

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

Upvotes: 2

gsmafra
gsmafra

Reputation: 2504

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

Upvotes: 4

Related Questions