num/argmax.go at master · quant1x/num · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
package num

import (
	"gitee.com/quant1x/num/x32"
	"gitee.com/quant1x/num/x64"
)

// ArgMax Returns the indices of the maximum values along an axis.
//
//	返回轴上最大值的索引
func ArgMax[T Number](x []T) int {
	ret := UnaryOperations2[T, int](x, x32.ArgMax, x64.ArgMax, __go_arg_max[T])
	return ret
}

func ArgMax2[T BaseType](x []T) int {
	var d int
	switch vs := any(x).(type) {
	case []float32:
		d = ArgMax(vs)
	case []float64:
		d = ArgMax(vs)
	case []int:
		d = ArgMax(vs)
	case []int8:
		d = ArgMax(vs)
	case []int16:
		d = ArgMax(vs)
	case []int32:
		d = ArgMax(vs)
	case []int64:
		d = ArgMax(vs)
	case []uint:
		d = ArgMax(vs)
	case []uint8:
		d = ArgMax(vs)
	case []uint16:
		d = ArgMax(vs)
	case []uint32:
		d = ArgMax(vs)
	case []uint64:
		d = ArgMax(vs)
	case []uintptr:
		d = ArgMax(vs)
	case []string:
		d = __go_arg_max(vs)
	case []bool:
		d = __go_bool_arg_max(vs)
	default:
		// 其它类型原样返回
		panic(TypeError(any(x)))
	}

	return d
}

func __go_arg_max[T Ordered](x []T) int {
	maxValue := x[0]
	idx := 0
	for i, v := range x[1:] {
		if v > maxValue {
			maxValue = v
			idx = 1 + i
		}
	}
	return idx
}

func __go_bool_arg_max(x []bool) int {
	maxValue := BoolToInt(x[0])
	idx := 0
	for i, v := range x[1:] {
		if BoolToInt(v) > maxValue {
			maxValue = BoolToInt(v)
			idx = 1 + i
		}
	}
	return idx
}