CODEWITHSUNDEEP
pythonnumpynumerical-methods
linello
linello

Reputation: 8694

Triangular indices for multidimensional arrays in numpy

We know that np.triu_indices returns the indices of the triangular upper part of a matrix, an array with two dimensions.

What if one wants to create indices as in the following code?

indices = []
for i in range(0,n):
    for j in range(i+1,n):
        for k in range(j+1,n):
            indices.append([i,j,k])

in a num-pythonic way?

Upvotes: 2

Views: 1635

Answers (1)

mgab
mgab

Reputation: 3964

In general you can get a list of indexes that follow the logic of the code you put with

from itertools import combinations
ndim = 3 # number of dimensions
n = 5 # dimension's length (assuming equal length in each dimension)

indices = list(combinations(range(n), r=ndim)

or if you want to iterate over each position:

for i,j,k in combinations(range(n), r=ndim):
    # Do your cool stuff here
    pass

However, you referred to it as the triangular upper part of a multidimensional matrix. I'm not sure what's the definition of it, and trying to visualize the indexes you selected with your nested loops I can't figure it out... (I'm just curious now on if there's a deffinition for multidimensional triangular matrix :-P)

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

a = zip(*indices)
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(a[0], a[1], a[2])
plt.xlabel('x')
plt.ylabel('y')
plt.show()

enter image description here (I moved the view angle to try to show what positions are selected)

Upvotes: 3

Related Questions