Getting indices of different lengths to slice a multidimensional numpy array

Question

Given these arrays,

  R = np.array([[ 5.,  3.,  2.,  7.,  3.,  6.,  8.,  9., 10., 55.],
                  [ 5.,  4.,  2.,  7.,  3.,  6.,  8., 10., 10., 55.]])
  F = np.array([[ 0.2 ,  0.4 ,  0.1 ,  0.3 ,  0.25,  0.25,  0.2 ,  0.1 ,  0.1 , 0.1 ],
                  [ 0.3 , -0.4 ,  0.1 ,  0.3 ,  0.25,  0.25,  0.4 , -0.4 ,  0.1 , 0.1 ]])
  K = nparray([[2],
                 [1]])

I want to sort each row of F and then find the first K[0] indices of the first row and the first K[1] indices of the second row of the sorted F. Then using these indices, I want to add 1 to the elements of R.

Here is my attempt, I was able to get the indices:

indices = np.argsort(F)[np.tile(np.arange(F.shape[1]),(F.shape[0],1)) < K]
# indices =  np.array([7, 8, 7], dtype=int64)

But I am not sure how to vectorize the below operations without using a for loop.

Rnew = R.copy()

Rnew[0,7] =R[0,7]+1

Rnew[0,8] =R[0,8]+1

Rnew[1,7] =R[1,7]+1

Rnew = np.array([[ 5.,  3.,  2.,  7.,  3.,  6.,  8., 10., 11., 55.],
                 [ 5.,  4.,  2.,  7.,  3.,  6.,  8., 11., 10., 55.]])

bb1 · Accepted Answer · 2021-08-20T02:27:19.757

I think this should work:

m = np.arange(F.shape[1]) < K
Rnew = R.copy()
Rnew[np.nonzero(m)[0], np.argsort(F)[m]] += 1

Since the first line uses broadcasting, np.tile() is not needed.

Notice that there is a possible ambiguity of the results: since each row of F has values that are repeated several times (e.g. 0.1 in the first row and -0.4 in the second) np.argsort() may give different orderings of elements of F, depending on how these equal values get sorted. This may change which entries of the matrix R will get incremented. For example, instead of incrementing R[0, 7], R[0, 8], and R[1, 7], the code may increment the entries R[0, 2], R[0, 9] and R[1, 1]. To get unambiguous results you can specify that np.argsort() must use a stable sorting algorithm, which will preserve the relative order of elements with equal values:

m = np.arange(F.shape[1]) < K
Rnew = R.copy()
Rnew[np.nonzero(m)[0], np.argsort(F, kind="stable")[m]] += 1

In this particular example this will increment the entries R[0, 2], R[0, 7] and R[1, 1]. You need to decide if this is the result that meets your needs.

Getting indices of different lengths to slice a multidimensional numpy array

1 Answers1