How to randomly select an element in a 2D numpy array in the same row or column of specified element using vectorization?

user_3337089 0
Viewed 111

I have a 2D NumPy array (say arr1) containing values 0 or 1 as float values. Let size of arr1 be h x w. I have another NumPy array (say arr2) of size n x 2, where each row specifies a location (row and column index) in arr1. For every arr1 location (say (x1, y1) ) specified by each row of arr2, I need to select another location (say (x2, y2)) in arr1 which is in the same row or column as (x1, y1), such that there is atleast one cell between (x1, y1) and (x2, y2), including these two cells, whose value in arr1 is 1. How can I achieve this efficiently in time? Typical values of h,w,n are 800,800,500000 respectively. So, I would like to achieve this without any for loops.

Example:

import numpy

h=4
w=4
n=3
arr1 = numpy.array([
        [0, 1, 0, 0],
        [1, 0, 1, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
      ])
arr2 = numpy.array([
        [1, 1],
        [2, 2],
        [0, 2],
      ])

Expected solution: First row of arr2 is (1,1). Valid solutions in 2nd column as (0,1), (2,1), (3,1) and valid solutions in 2nd row are (1,0), (1,2), (1,3). So the code should randomly pick one of these. Similar for second row of arr2 which is (2,2), valid solutions are (0,2), (1,2), (3,2), (2,0), (2,1), (2,3). For third row of arr2 which is (0,2), valid solutions are (0,0),(0,1),(1,2),(2,2),(3,2). Note that (0,3) is not a valid solution since there is no cell containing 1 between (0,2) and (0,3).

Note that if a row in arr2 is (0,3), there is no cell in that column with the value 1. Such cases are extremely rare and in such cases, it suffices to pick a location that is sufficiently far away in that column. It is not necessary to detect such cases and pick a location in the same row.

PS: I have a solution by iterating over each row of arr2, but that takes over 1 minute. I am looking for a vectorized solution

vectorization numpy python
1 Answers

I tried writing a solution to this using vectorization, but it still need to uses masked reduce operations which makes things way slower than normal broadcasting

The basic idea behind what i did is to rotate every row by the col index in arr2 and find the min/max index of a non 0 element. Than you select a random number between this 2 (all possible rotated solutions) and rotate it back to get the actual solution for that row. You than do the same on the transposed arr1 and arr2 to get the the solution by columns and than you randomly select if you want the row or colum solution.

Still this take ~26s to execute

PS: If no solution is possible for the initial point (both row and col are all zeros), a -1 will appear in those coordinates

import time
import numpy as np

def get_same_row_random(arr, pos, shift):
    width = arr.shape[1]
    app1 = -np.ones_like(arr, dtype=np.int64)
    app2 = np.zeros_like(arr, dtype=np.bool8)
    w1 = np.where(arr == 1)
    app1[w1] = w1[1]
    app2[w1] = True

    res = app1[pos]
    mask = app2[pos]

    mask[np.arange(shift.size), shift] = False # Can't select original coordinates
    res = (res - shift.reshape(-1,1)) % width
    res_min = np.min(res, initial=width, where=mask, axis=1)
    res_max = np.max(res, initial=-1, where=mask, axis=1)

    w = (res_min == width)
    res_min[w] = -1
    res_max += 1

    res = (np.random.randint(res_min, res_max) + shift) % width
    res[w] = -1

    return res

h, w, n = 800,800, 500000
arr1 = np.random.randint(2, size=(h,w))
arr2_1 = np.random.randint(h, size=n)
arr2_2 = np.random.randint(w, size=n)
arr2 = np.vstack((arr2_1, arr2_2)).T

start = time.time()

row_indexes = get_same_row_random(arr1, arr2[:,0], arr2[:,1])
col_indexes = get_same_row_random(arr1.T, arr2[:,1], arr2[:,0])

r_or_c = np.random.randint(2, size=n)
exc_row = np.where(row_indexes == -1)[0]
exc_col = np.where(col_indexes == -1)[0]
r_or_c[exc_row] = 1
r_or_c[exc_col] = 0

wr = np.where(r_or_c == 0)[0]
wc = np.where(r_or_c == 1)[0]
res = np.array(arr2)
res[wr,1] = row_indexes[wr]
res[wc,0] = col_indexes[wc]

end = time.time()
user_76044340
Related