CODEWITHSUNDEEP
pythonnumpy
ChristosK
ChristosK

Reputation: 25

Get indices from np.array under condition from another array

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".

Upvotes: 1

Views: 128

Answers (1)

Tonechas
Tonechas

Reputation: 13733

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)

Upvotes: 1

Related Questions