0

I have a matrix

A = np.zeros((n,n))

with say n=4. I have another matrix

B = np.array(
    [7, 3, 5, 4],
    [4, 3, 2, 1],
    [6, 7, 4, 5],
    [1, 2, 3, 4]
)

First I need the indices of the smallest k elements (e.g. 2) in every row of B, in this case it would be

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

I can do this doing np.argsort but that ends up sorting the whole array, np.argpartition can sort the first k elements but I need a way to get the indices. Lastly I want to insert the smallest k values of B into A at the indices giving the matrix

A = np.array(
    [0, 3, 0, 4],
    [0, 0, 2, 1],
    [0, 0, 4, 5],
    [1, 2, 0, 0]
)

What's a possible way to this?

mkrieger1
  • 19,194
  • 5
  • 54
  • 65
John
  • 81
  • 4
  • Does this answer your question? [Find the index of the k smallest values of a numpy array](https://stackoverflow.com/questions/34226400/find-the-index-of-the-k-smallest-values-of-a-numpy-array) – mkrieger1 Apr 04 '21 at 21:58
  • It answers the first part of the question about getting the indices of the k smallest values, I need the last part which inserts the values at those indices in `B` into `A`. – John Apr 04 '21 at 22:04
  • Have you tried `A[indices] = B[indices]`? I'm not sure but it may be as simple as that. – mkrieger1 Apr 04 '21 at 22:04

1 Answers1

0

Try this:

n = 4
A = np.zeros((n,n), dtype=np.int32)
B = np.array([[7,3,5,4],[4,3,2,1],[6,7,4,5],[1,2,3,4]])

indices = np.argpartition(B, 2, axis=1)[:,:2]
values = np.take_along_axis(B, indices, axis=1)
np.put_along_axis(A, indices, values, axis=1)
print(A)

Output:

array([[0, 3, 0, 4],
       [0, 0, 2, 1],
       [0, 0, 4, 5],
       [1, 2, 0, 0]], dtype=int32)
Kevin
  • 3,096
  • 2
  • 8
  • 37