[Numpy-discussion] argmax() indexes to value

[CJ Carey]

You could move some of the cost to index-creation time by converting the
per-row indices into flattened indices:

In [1]: a = np.random.random((5, 6))


In [2]: i = a.argmax(axis=1)


In [3]: a[np.arange(len(a)), i]

Out[3]: array([0.95774465, 0.90940106, 0.98025448, 0.97836906, 0.80483784])

In [4]: f = np.ravel_multi_index((np.arange(len(a)), i), a.shape)


In [5]: a.flat[f]

Out[5]: array([0.95774465, 0.90940106, 0.98025448, 0.97836906, 0.80483784])


I haven't benchmarked, but I suspect this will be faster if you're using
the same index multiple times.
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20191101/b492808f/attachment.html>