0

numpy 配列があり、最初の軸を除くすべての軸で argmax を取得したいと思います。私は解決策を持っています（と思います）が、それを行うためのより良い方法があるかどうか疑問に思います.

import numpy as np

def argmax(array):
    ## Argmax along all axes except the first (ie axis 0)
    last_axis = len(array.shape) - 1
    if last_axis == 0:
        return tuple(range(array.size))
    if last_axis == 1:
        return (range(array.shape[0]), list(np.argmax(array, axis=1)))
    index_array = np.argmax(array, axis=last_axis)
    smaller_array = np.amax(array, axis=last_axis)
    assert index_array.shape == smaller_array.shape
    argmax_smaller_array = argmax(smaller_array)
    return argmax_smaller_array + (list(index_array[argmax_smaller_array]), )

いくつかの例：

a = np.arange(12).reshape((6, 2))
a[5, 0] = 22
argmax(a)
a[argmax(a)]

b = np.arange(18).reshape((3, 3, 2))
b[0, 0, 0] = 55
b[argmax(b)]
np.all(b[argmax(b)] == np.array([np.max(b[0]), np.max(b[1]), np.max(b[2])]))  # True

私は numpy を使い始めたばかりで、これを行う簡単な方法があるかどうか疑問に思っています。私はすでに存在するものを書き直していますか？

4

1 に答える 1

2

あなたの方法は問題ないようですが、必要のない多くの中間結果を計算しています。次のようなことができます：

import numpy as np

def argmax(array):
    shape = array.shape
    array = array.reshape((shape[0], -1))
    ravelmax = np.argmax(array, axis=1)
    return (np.arange(shape[0]),) + np.unravel_index(ravelmax, shape[1:])
于 2014-03-21T17:53:19.260 に答える