Numpy Indices of a masked array with multiple conditions-CodePudding

What is the most efficient way to get the indexes of the upper right triangle with values of 1?

Given:

A = np.array([
    [0,1,0,1],
    [1,0,1,0],
    [0,1,0,0],
    [1,0,0,0]])

I want:

[(0,1), (0,3), (1,2)]

np.where(A[np.triu_indices(A.shape[0])]==1) gives me it as a 1D array and not as a tuple. Looked at using np.unravel_index(mask,(4,4)) but that gave:

(array([[0, 0, 1]], dtype=int64), array([[1, 3, 1]], dtype=int64))

Which is not correct as the last tuple should be 1,2?

CodePudding user response：

You can use np.triu to make a mask, np.where to get the indexes, and zip list to get them into the right format:

indexes = list(zip(*np.where(np.triu(A))))

Output:

>>> indexes
[(0, 1), (0, 3), (1, 2)]

CodePudding user response：

In [77]: A = np.array([[0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]])

The upper tri:

In [78]: np.triu(A)
Out[78]: 
array([[0, 1, 0, 1],
       [0, 0, 1, 0],
       [0, 0, 0, 0],
       [0, 0, 0, 0]])

The nonzero indices - as a tuple of arrays that can be used to index A:

In [79]: np.nonzero(np.triu(A))
Out[79]: (array([0, 0, 1]), array([1, 3, 2]))

the transpose of that gives an array of the index pairs. It's closer to what you want, though not as convenient for indexing:

In [80]: np.argwhere(np.triu(A))
Out[80]: 
array([[0, 1],
       [0, 3],
       [1, 2]])

Index with the tuple vs with the pairs:

In [84]: A[np.nonzero(np.triu(A))]
Out[84]: array([1, 1, 1])
In [85]: [A[tuple(pair)] for pair in np.argwhere(np.triu(A))]
Out[85]: [1, 1, 1]