Distribute elements across given number of buckets (evenly and based on weights)

bucketise.go

package main

import (
    "encoding/csv"
    "fmt"
    "math/rand"
    "os"
    "strconv"
)

func main() {
    n, _ := strconv.Atoi(os.Args[1])

    seq := make([]int, n)
    for i := 0; i < n; i++ {
        seq[i] = i + 1
    }

    buckets := bucketiseBiased(seq, []float64{1.0 / 6.0, 3.0 / 6.0, 2.0 / 6.0})
    for i, bucket := range buckets {
        fmt.Printf("Bucket #%d has %d elements\n", i+1, len(bucket))
    }

    fmt.Println("---")

    buckets = bucketiseEven(seq, 3)
    for i, bucket := range buckets {
        fmt.Printf("Bucket #%d has %d elements\n", i+1, len(bucket))
    }

    fmt.Println("---")

    res := make(map[string]int)
    data := []int{1, 2, 3, 4}
    for i := 0; i < 600000; i++ {
        shuffleKFY(data)
        res[fmt.Sprintf("%#v", data)]++
    }
    csv := csv.NewWriter(os.Stdout)
    for k, v := range res {
        csv.Write([]string{k, strconv.Itoa(v)})
    }
    csv.Flush()
}

func bucketiseBiased(data []int, weights []float64) [][]int {
    buckets := make([][]int, len(weights))

    for i, _ := range data {
        bucket := biasedDice(weights)
        buckets[bucket] = append(buckets[bucket], data[i])
    }

    return buckets
}

func bucketiseEven(data []int, bucketsNum int) [][]int {
    buckets := make([][]int, bucketsNum)

    shuffleKFY(data)

    minBucketSize := len(data) / bucketsNum
    for bucket := 0; ; bucket++ {
        if len(data) < minBucketSize {
            buckets[bucket-1] = append(buckets[bucket-1], data...)
            break
        } else {
            buckets[bucket] = data[:minBucketSize]
            data = data[minBucketSize:]
        }
    }

    return buckets
}

func shuffleKFY(data []int) {
    for i := len(data) - 1; i > 0; i-- {
        j := rand.Intn(i + 1)
        data[i], data[j] = data[j], data[i]
    }
}

func bucketiseEvenRR(data []int, bucketsNum int) [][]int {
    buckets := make([][]int, bucketsNum)

    i := 0
    total := len(data)
    for i < total {	
	    idx := rand.Intn(len(data))
	    elem := data[idx]
	    data = append(data[:idx], data[idx+1:]...)

	    currBucket := i % bucketsNum
	    i++
	    buckets[currBucket] = append(buckets[currBucket], elem)
    }

    return buckets
}

func biasedDice(weights []float64) int {
    var sum float64
    for _, v := range weights {
        sum += v
    }

    r := rand.Float64() / float64(sum)
    for i, v := range weights {
        r -= v
        if r < 0 {
            return i
        }
    }

    panic("should not happen")
}