Go語言：按照指定的一組權重隨機選擇一個數組元素（加權隨機索引）

加權隨機索引比較常用，例如HTTP代理服務器需要按照不同的權重把請求轉發給不同的後端服務器，又比如在web頁面的同一個位置需要按照一定的比例顯示不同的廣告圖片。

函數實現如下：

package main

import (
	"fmt"
	"math/rand"
	"time"
)

/*
  作者：pengpengzhou
  方法：WeightedRandomIndex
  功能：按照指定的一組權重隨機返回數組索引
  參數：weights []float32 權重切片
  返回：加權隨機索引index，index是 0 ~ len(weights)-1 之間的一個整數

  示例如下：
  按權重[0.1, 0.2, 0.3, 0.4]隨機調用1000次該方法，返回0,1,2,3的次數將接近於1:2:3:4

  	var weights = []float32{0.1, 0.2, 0.3, 0.4}
  	var result [4]int
  	rand.Seed(time.Now().Unix())
  	for i := 0; i < 1000; i++ {
  		result[WeightedRandomIndex(weights)]++
  	}
	fmt.Printf("%v\n", result)

  輸出：
    [112 174 304 410]
*/
func WeightedRandomIndex(weights []float32) int {
	if len(weights) == 1 {
		return 0
	}
	var sum float32 = 0.0
	for _, w := range weights {
		sum += w
	}
	r := rand.Float32() * sum
	var t float32 = 0.0
	for i, w := range weights {
		t += w
		if t > r {
			return i
		}
	}
	return len(weights) - 1
}

func main() {
	var weights = []float32{0.1, 0.2, 0.3, 0.4}
	var result [4]int
	rand.Seed(time.Now().Unix())
	for i := 0; i < 1000; i++ {
		result[WeightedRandomIndex(weights)]++
	}
	fmt.Printf("%v\n", result)
}

輸出：

[108 189 291 412]

以Go內置的rand.Intn函數作爲參照，進行性能對照測試：

package main

import (
	"math/rand"
	"testing"
	"time"
)

var weights = []float32{0.1, 0.2, 0.3, 0.4}
var n int = len(weights)

func BenchmarkWeightedRandomIndex(b *testing.B) {
	rand.Seed(time.Now().Unix())
	for i := 0; i < b.N; i++ {
		WeightedRandomIndex(weights)
	}
}

func BenchmarkRandomIntn(b *testing.B) {
	rand.Seed(time.Now().Unix())
	for i := 0; i < b.N; i++ {
		rand.Intn(n)
	}
}

測試結果：

[root@dev example]# go test -bench=.
goos: linux
goarch: amd64
pkg: example
BenchmarkWeightedRandomIndex-4          34809902                34.1 ns/op
BenchmarkRandomIntn-4                   51488966                23.3 ns/op
PASS
ok      example 3.340s

可以看出：加權隨機索引函數比等概率隨機索引函數執行效率略慢，但大致相當。