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.
_______________________________________________ NumPy-Discussion mailing list NumPy-Discussion@python.org https://mail.python.org/mailman/listinfo/numpy-discussion
