[SOLVED] Get indices from np.array under condition from another array

Issue

Suppose I have a NumPy array (5×5) with distances between points.
The matrix is square and symmetrical.

distances = np.array([
    [0, 3, 2, 1, 4],
    [3, 0, 1, 3, 5],
    [2, 1, 0, 7, 6],
    [1, 3, 7, 0, 9],
    [4, 5, 6, 9, 0],
    ])

I also have a points list of 5 elements as following:

points = ["A", "D", "A", "D", "F"]

so that, for example, the distance between the first "A" and the first "D" of the points list is distances[0, 1], which is 3.

Is it possible, using the np.where() function to get the indices where the distance is less than 2, the first point is "A" and the second point is "D"? That is, to get [0, 4] and [2, 1]?

If I use np.where((distances <= 2) & (distances > 0)) I get the indices of the array for all values <=2. I want to get the indices for those values <= 2, where the first point is "A" and the second point is "D".

Solution

Something like this?

In [152]: import numpy as np

In [153]: points = np.array(['A', 'D', 'A', 'D', 'F'], dtype='str')

In [154]: distances = np.array(
     ...:     [[0, 3, 2, 1, 4],
     ...:      [3, 0, 1, 3, 5],
     ...:      [2, 1, 0, 7, 6],
     ...:      [1, 3, 7, 0, 9],
     ...:      [4, 5, 6, 9, 0]])

In [155]: first = 'A'

In [156]: second = 'D'

In [157]: threshold = 2

In [158]: rows = np.flatnonzero(points == first)

In [159]: cols = np.flatnonzero(points == second)

In [160]: ixgrid = np.ix_(rows, cols)

In [161]: idx1, idx2 = np.where(distances[ixgrid] <= threshold)

In [162]: np.stack([rows[idx1], cols[idx2]]).T
Out[162]: 
array([[0, 3],
       [2, 1]], dtype=int64)

Answered By – Tonechas

Answer Checked By – David Marino (BugsFixing Volunteer)

Leave a Reply

Your email address will not be published. Required fields are marked *