Is there a vectorized way of appending/indexing values from nd_array 2 based off component r/c value-CodePudding

I'm looking to vectorize some code that I have, but am unsure how to approach it or if it's even possible. Here's what I have so far:

arr_row = [0, 1, 2, ..., n] # Array of size n with random integers
arr_col = [0, 1, 2, ..., n] # Array of size n with random integers
for i in range(n):
     r = arr_row[i]
     c = arr_col[i]
     result = n_by_n_table[r,c]

So briefly about the above. I have two arrays that are n elements, I iterate through i to n getting the item at the row/col location obtained from my row/col arrs. I get this result from a guaranteed n x n matrix. I am hoping that I could vectorized this for a performance improvement. Here's what I had in mind:

vectorized = nd_array((n, 3))
vectorized[:, 0] = [0, 1, 2, ..., n] # Likely using np.rand func or something here.
vectorized[:, 1] = [0, 1, 2, ..., n] # Same as above
vectorized[:, 2] = n_by_n_table[vectorized[:,0], vectorized[:,1]]

Is this possible and any concerns doing the above?

CodePudding user response：

example data

arr_row = [5,1,0,3,2,4]
arr_col = [1,4,2,0,5,3]
n_by_n_table = np.array(range(36)).reshape(6,6)

n_by_n_table :

[[ 0,  1,  2,  3,  4,  5],
 [ 6,  7,  8,  9, 10, 11],
 [12, 13, 14, 15, 16, 17],
 [18, 19, 20, 21, 22, 23],
 [24, 25, 26, 27, 28, 29],
 [30, 31, 32, 33, 34, 35]]

solution

n_by_n_table[arr_row,arr_col]

result:

array([31, 10,  2, 18, 17, 27])