Numpy配列で指定の軸の2番目に大きい値とインデックスを求める

Numpy配列で、指定の軸の1番大きい値とインデックスは、numpy.amax と numpy.argmax で求まる。では、2 番目に大きい値とインデックスはどうするか？

以下の配列で、それぞれの行 (axis=1) の場合を考える。

import numpy as np

a = np.array([[20, 3, 100, 4000], [1, 400, 2000, 30], [300, 1000, 40, 2]])
print(a)
# -----
[[  20    3  100 4000]
 [   1  400 2000   30]
 [ 300 1000   40    2]]

まず、numpy.argsort で、指定の軸に沿ってソートすると、それぞれの値が何番目になるかを求める。

a_argsorted = np.argsort(a, axis=1)
print(a_argsorted)
# -----
[[1 0 2 3]
 [0 3 1 2]
 [3 2 0 1]]

最初の行 [ 20 3 100 4000] の結果は [1 0 2 3] で、一番小さいのがインデックス 1 の 3、2 番目に小さいのがインデックス 0 の 20、その次はインデックス 2 の 100、最後がインデックス 3 の 4000 ということになる。それぞれの行で最大の値のインデックスは、argsort した結果の各行の最後の値、2 番目に大きい値のインデックスは、各行の後ろから 2 番目の値になる。

a_2nd_largest_idx = a_argsorted[ : , -2 ]
print(a_2nd_largest_idx)
# -----
[2 1 0]

各行で 2 番目に大きい値は、各行のこのインデックスの値を取り出せばよい。

a_2nd_largest_val = a[ np.arange(a.shape[0]) , a_2nd_largest_idx ]
print(a_2nd_largest_val)
# -----
[100 400 300]