CODEWITHSUNDEEP
pythonnumpy
exception1
exception1

Reputation: 1249

Numpy: Column dependent slicing

I could not find any question concerning what I want to do so I am asking now. Basically, I want slicing in matrices where the row index depends on the column index.

For example:

>>> import numpy as np
>>> x = np.arange(24).reshape(6,4)
>>> x
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15],
       [16, 17, 18, 19],
       [20, 21, 22, 23]])

Now what I need is:

array([[12,  9,  6,  3],
       [16, 13, 10,  7],
       [20, 17, 14, 11]])

So I have a given (fixed) set of column indices. The row indices depend on them. For example: row = col-1:col+1

Is there an efficient way to do that? I know I can iterate and slice the entries I need from every single column, but it seems very inefficient to me.

Upvotes: 0

Views: 131

Answers (2)

hpaulj
hpaulj

Reputation: 231335

I was going to leave a diagonal based solution in the comments, but in time tests it proves to be faster than the strided one. 

In [128]: timeit d = as_strided(x[::-1,:], shape=(3,4),
     strides=sz*np.array([-4,-3]))[::-1,:]
10000 loops, best of 3: 44 µs per loop

In [129]: timeit np.array([x[::-1,:].diagonal(i) for i in [-2,-1,0]])
10000 loops, best of 3: 25.9 µs per loop

With a larger x this probably will not be true. But the diagonal solution has the advantage that it might be easier to understand.

for a (69,40) x, 

np.array([x[::-1,:].diagonal(i) for i in range(-(n-1),1)])

produces a (30,40) array that includes all of the length 40 reverse diagonals.

The as_strided solution is:

as_strided(x[::-1,:], shape=(n,m), strides=-sz*np.array([m,m-1]))[::-1,:]

It wasn't as easy to figure out how to adjust the values to produce the correct strides. But it is faster. Time basically the same as for the smaller array, while the diagonal approach slows down with size.

Upvotes: 2

xnx
xnx

Reputation: 25478

You can use as_strided to do this:

In [1]: from numpy.lib.stride_tricks import as_strided
In [2]: sz = x.itemsize
In [3]: d = as_strided(x[-1::-1,:], shape=(3,4), strides=sz*np.array([-4,-3]))
In [4]: d
Out[5]: 
array([[20, 17, 14, 11],
       [16, 13, 10,  7],
       [12,  9,  6,  3]])

That is, starting with the original array upside-down, fill a 3x4 array with items taken by striding backwards 4 items for the first coordinate and 3 for the second.

If you want, you can then view d upside-down:

In [6]: d[-1::-1,:]    # or np.flipud(d)
Out[6]: 
array([[12,  9,  6,  3],
       [16, 13, 10,  7],
       [20, 17, 14, 11]])

Upvotes: 4

Related Questions