arraysnumpy

Efficiently identifying instances in 2d arrays where column elements match two conditions


Consider the 3d array arr, shown below. I need to find instances in arr[0], arr[1], and arr[2] where the following condition holds: the 2d array holds element p in row q and element q in row p, both in the SAME COLUMN. For example, column 2 of arr[0] holds a 3 in row 2 and a 2 in row 3. There are other instances as well.

The following code identifies these instances, but it has multiply nested for and if statements, and this will be slow for large numbers of larger arrays. QUESTION: is there a more efficient way to do this?

import numpy as np

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

for i in range(3):     
    curr_arr = arr[i]
    num_rows, num_cols = curr_arr.shape
    for col in range(num_cols):
        for row1 in range(num_rows):
            for row2 in range(num_rows):
                if row1 != row2 and row1 < row2:
                    
                    val1 = curr_arr[row1, col]
                    val2 = curr_arr[row2, col]
                    
                    if val1==row2 and val2 == row1:
                        print(i)
                        print(row1, val1)
                        print(row2, val2)                 
                        print()

Ideally, the output would show i, row1, val1, row2, val2 in a numpy 2d array. The solutions to this example would be:

output = [[0, 0, 3, 3, 0],
          [0, 1, 2, 2, 1],
          [0, 2, 3, 3, 2],
          [1, 0, 2, 2, 0],
          [1, 2, 3, 3, 2],
          [2, 1, 3, 3, 1],
          [2, 0, 3, 3, 0],
          [2, 1, 2, 2, 1]]

Thanks for any assistance.


Solution

  • A possible solution:

    k, m, n = arr.shape
    results = []
    for i in range(k):
        for col in range(n):
            col_vec = arr[i, :, col]
            for p in range(m):
                q = col_vec[p]
                if q < 0 or q >= m:
                    continue
                if p < q and col_vec[q] == p:
                    results.append([i, p, q, q, p])
    
    np.array(results)
    

    To find symmetric row pairs (p, q) in each 2D slice of a 3D array where arr[i, p, col] = q and arr[i, q, col] = p (with p < q), it leverages direct indexing to avoid nested loops. For each slice i, column col, and row p: (i) it accesses the column vector using col_vector = arr[i, :, col]; (ii) for each p, get q = col_vector[p]. If q is a valid row index (0 ≤ q < m) and p < q, proceed; and (iii) check if col_vector[q] == p, and, if true, record [i, p, q, q, p].

    Output:

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