numpy - Kotlin 是否有类似于 argmax 方法的东西？

Question

所以假设我有一个像这样的numpy数组：

import numpy as np
mat = np.array([[4, 8, 1], [5, 10, 6]])

print(np.argmax(mat)) # prints 4
print(np.argmax(mat, axis=1)) # prints [1 1], index of maximum values along the rows

Kotlin 是否有类似的（内置）功能？我找到了NumPy 的 Kotlin 绑定，但我没有找到实现的功能。

提前致谢！

Answer 1

使用withIndex()和maxByOrNull()：

fun <T : Comparable<T>> Iterable<T>.argmax(): Int? {
    return withIndex().maxByOrNull { it.value }?.index
}