gitee.com/quant1x/num@v0.3.2/where.go (about)

     1  package num
     2  
     3  // Where 返回根据“条件”从“x”或“y”中选择的元素
     4  // 这里先实现一个简单的, 留给于总重构
     5  // params只支持两个默认值x和y, 如果condition为true返回x, 否则返回y
     6  // condition和param都可能是基础数据类型,也可能是一个slice, 并且长度可能不一致
     7  // 直接写成序列版本, 可能更简单
     8  func Where[T StatType](condition []T, x, y []T) []T {
     9  	// 第一步, 找出最长的
    10  	clen := len(condition)
    11  	xlen := len(x)
    12  	ylen := len(y)
    13  	// 第二步, 找出最大长度
    14  	c := []float64{float64(clen), float64(xlen), float64(ylen)}
    15  	maxLength := int(Max(c))
    16  
    17  	// 处理默认值
    18  	defaultValue := TypeDefault[T]()
    19  	// 对齐所有长度
    20  	if clen < maxLength {
    21  		condition = Align(condition, defaultValue, maxLength)
    22  	}
    23  	if xlen < maxLength {
    24  		x = Align(x, defaultValue, maxLength)
    25  	}
    26  	if ylen < maxLength {
    27  		y = Align(y, defaultValue, maxLength)
    28  	}
    29  	// 初始化返回值
    30  	d := make([]T, maxLength)
    31  	for i := 0; i < maxLength; i++ {
    32  		// NaN也被认为是真
    33  		if condition[i] != 0 {
    34  			d[i] = x[i]
    35  		} else {
    36  			d[i] = y[i]
    37  		}
    38  	}
    39  	return d
    40  }