Numpy查找多维矩阵第k大值的位置（索引） | Using Numpy to find the index of the k-th maximum in a multidimensional matrix

发表于 | 分类于

Find the index of maximum

I tried hard to find the index of maximum in a multidimensional matrix, and here is the code:

1
np.unravel_index(np.argmax(MATRIX), MATRIX.shape)

where argmax returns the index of the maximum in the flattened array of MATRIX, unravel_index transforms this index into a standard form of index.

Find the index of K-th Maximum

Same as above, the code is:

1
2
flattened_index = np.argpartition(M.flatten(), -K)[-K]
index = np.unravel_index(flattened_index, M.shape)

For example, matrix a is:

1
2
3
0	0	0	0	0
0	0	0	1	2
0	0	0	0	0

Then:

1
2
3
4
5
flattened_index = np.argpartition(a.flatten(), -2)[-2]
index = np.unravel_index(flattened_index, a.shape)

assert flattend_index == 8
assert index == (1, 3)
Tipping