2Cubed - 3 months ago 11
Python Question

# Find all n-dimensional lines and diagonals with NumPy

Using NumPy, I would like to produce a list of all lines and diagonals of an n-dimensional array with lengths of k.

Take the case of the following three-dimensional array with lengths of three.

``````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],
[24, 25, 26]]])
``````

For this case, I would like to obtain all of the following types of sequences. For any given case, I would like to obtain all of the possible sequences of each type. Examples of desired sequences are given in parentheses below, for each case.

• 1D lines

• x axis (
`0, 1, 2`
)

• y axis (
`0, 3, 6`
)

• z axis (
`0, 9, 18`
)

• 2D diagonals

• x/y axes (
`0, 4, 8`
,
`2, 4, 6`
)

• x/z axes (
`0, 10, 20`
,
`2, 10, 18`
)

• y/z axes (
`0, 12, 24`
,
`6, 12, 18`
)

• 3D diagonals

• x/y/z axes (
`0, 13, 26`
,
`2, 13, 24`
)

The solution should be generalized, so that it will generate all lines and diagonals for an array, regardless of the array's number of dimensions or length (which is constant across all dimensions).

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 `bi`s come before `fi`s, 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
sl = np.index_exp[0]
for i in range(n):
sl = np.index_exp[2:] + sl

``````

``````# construct your example array
n = 3
k = 3
data = np.arange(k**n).reshape((k,)*n)

# apply my index_creating function

# apply the index to your array
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