Find all n-dimensional lines and diagonals with NumPy

This solution generalized over n

Lets rephrase this problem as "find the list of indices".

We're looking for all of the 2d index arrays of the form

array[i[0], i[1], i[2], ..., i[n-1]]

Let n = arr.ndim

Where i is an array of shape (n, k)

Each of i[j] can be one of:

• The same index repeated n times, ri[j] = [j, ..., j]
• The forward sequence, fi = [0, 1, ..., k-1]
• The backward sequence, bi = [k-1, ..., 1, 0]

With the requirements that each sequence is of the form ^(ri)*(fi)(fi|bi|ri)*$ (using regex to summarize it). This is because:

• there must be at least one fi so the "line" is not a point selected repeatedly
• no bis come before fis, to avoid getting reversed lines

def product_slices(n):    for i in range(n):        yield (            np.index_exp[np.newaxis] * i +            np.index_exp[:] +            np.index_exp[np.newaxis] * (n - i - 1)        )def get_lines(n, k):    """    Returns:        index (tuple):   an object suitable for advanced indexing to get all possible lines        mask (ndarray):  a boolean mask to apply to the result of the above    """    fi = np.arange(k)    bi = fi[::-1]    ri = fi[:,None].repeat(k, axis=1)    all_i = np.concatenate((fi[None], bi[None], ri), axis=0)    # inedx which look up every possible line, some of which are not valid    index = tuple(all_i[s] for s in product_slices(n))    # We incrementally allow lines that start with some number of `ri`s, and an `fi`    #  [0]  here means we chose fi for that index    #  [2:] here means we chose an ri for that index    mask = np.zeros((all_i.shape[0],)*n, dtype=np.bool)    sl = np.index_exp[0]    for i in range(n):        mask[sl] = True        sl = np.index_exp[2:] + sl    return index, mask

Applied to your example:

# construct your example arrayn = 3k = 3data = np.arange(k**n).reshape((k,)*n)# apply my index_creating functionindex, mask = get_lines(n, k)# apply the index to your arraylines = data[index][mask]print(lines)

array([[ 0, 13, 26],       [ 2, 13, 24],       [ 0, 12, 24],       [ 1, 13, 25],       [ 2, 14, 26],       [ 6, 13, 20],       [ 8, 13, 18],       [ 6, 12, 18],       [ 7, 13, 19],       [ 8, 14, 20],       [ 0, 10, 20],       [ 2, 10, 18],       [ 0,  9, 18],       [ 1, 10, 19],       [ 2, 11, 20],       [ 3, 13, 23],       [ 5, 13, 21],       [ 3, 12, 21],       [ 4, 13, 22],       [ 5, 14, 23],       [ 6, 16, 26],       [ 8, 16, 24],       [ 6, 15, 24],       [ 7, 16, 25],       [ 8, 17, 26],       [ 0,  4,  8],       [ 2,  4,  6],       [ 0,  3,  6],       [ 1,  4,  7],       [ 2,  5,  8],       [ 0,  1,  2],       [ 3,  4,  5],       [ 6,  7,  8],       [ 9, 13, 17],       [11, 13, 15],       [ 9, 12, 15],       [10, 13, 16],       [11, 14, 17],       [ 9, 10, 11],       [12, 13, 14],       [15, 16, 17],       [18, 22, 26],       [20, 22, 24],       [18, 21, 24],       [19, 22, 25],       [20, 23, 26],       [18, 19, 20],       [21, 22, 23],       [24, 25, 26]])

Another good set of test data is np.moveaxis(np.indices((k,)*n), 0, -1), which gives an array where every value is its own index

I've solved this problem before to implement a higher dimensional tic-tac-toe

In [1]: x=np.arange(27).reshape(3,3,3)

Selecting individual rows is easy:

In [2]: x[0,0,:]Out[2]: array([0, 1, 2])In [3]: x[0,:,0]Out[3]: array([0, 3, 6])In [4]: x[:,0,0]Out[4]: array([ 0,  9, 18])

You could iterate over dimensions with an index list:

In [10]: idx=[slice(None),0,0]In [11]: x[idx]Out[11]: array([ 0,  9, 18])In [12]: idx[2]+=1In [13]: x[idx]Out[13]: array([ 1, 10, 19])

Look at the code for np.apply_along_axis to see how it implements this sort of iteration.

Reshape and split can also produce a list of rows. For some dimensions this might require a transpose:

In [20]: np.split(x.reshape(x.shape[0],-1),9,axis=1)Out[20]: [array([[ 0],        [ 9],        [18]]), array([[ 1],        [10],        [19]]), array([[ 2],        [11],        ...

np.diag can get diagonals from 2d subarrays

In [21]: np.diag(x[0,:,:])Out[21]: array([0, 4, 8])In [22]: np.diag(x[1,:,:])Out[22]: array([ 9, 13, 17])In [23]: np.diag?In [24]: np.diag(x[1,:,:],1)Out[24]: array([10, 14])In [25]: np.diag(x[1,:,:],-1)Out[25]: array([12, 16])

And explore np.diagonal for direct application to the 3d. It's also easy to index the array directly, with range and arange, x[0,range(3),range(3)].

As far as I know there isn't a function to step through all these alternatives. Since dimensions of the returned arrays can differ, there's little point to producing such a function in compiled numpy code. So even if there was a function, it would step through the alternatives as I outlined.

==============

All the 1d lines

x.reshape(-1,3)x.transpose(0,2,1).reshape(-1,3)x.transpose(1,2,0).reshape(-1,3)

y/z diagonal and anti-diagonal

In [154]: i=np.arange(3)In [155]: j=np.arange(2,-1,-1)In [156]: np.concatenate((x[:,i,i],x[:,i,j]),axis=1)Out[156]: array([[ 0,  4,  8,  2,  4,  6],       [ 9, 13, 17, 11, 13, 15],       [18, 22, 26, 20, 22, 24]])

np.einsum can be used to build all these kind of expressions; for instance:

# 3d diagonalsprint(np.einsum('iii->i', a))# 2d diagonalsprint(np.einsum('iij->ij', a))print(np.einsum('iji->ij', a))