create_tree.go

package main

import (
	"fmt"
	"math"
	"sort"
)

const (
	maxLevel  = 10
	maxWeight = uint64(1) << 63
)

type (
	shiftDown struct {
		index int
		order float64
	}
	shiftList []shiftDown
)

func (s shiftList) Len() int           { return len(s) }
func (s shiftList) Less(i, j int) bool { return s[i].order < s[j].order }
func (s shiftList) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }

func calculateTreeLevels(values []float64) []int {
	var sum float64
	for _, v := range values {
		sum += v
	}
	mul := 0.99 / sum
	levels := make([]int, len(values))
	shift := make(shiftList, len(values))
	var totalWeight uint64

	for i, v := range values {
		l := -math.Log2(mul * v)
		c := math.Ceil(l)
		shift[i].index = i
		shift[i].order = l - c + 1
		level := int(c)
		if level <= maxLevel {
			totalWeight += maxWeight >> level
		} else {
			level = -1
		}
		levels[i] = level
	}

	sort.Sort(shift)

	for totalWeight < maxWeight && len(shift) != 0 {
		var j int
		for i, s := range shift {
			var addWeight uint64
			level := levels[s.index]
			if level > 0 {
				addWeight = maxWeight >> level
				level--
			} else {
				level = maxLevel
				addWeight = maxWeight >> maxLevel
			}
			if totalWeight+addWeight <= maxWeight {
				totalWeight += addWeight
				levels[s.index] = level
				if j != i {
					shift[j] = shift[i]
				}
				j++
			}
			if totalWeight == maxWeight {
				break
			}
		}
		shift = shift[:j]
	}
	return levels
}

func main() {
	fmt.Println(calculateTreeLevels([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9}))
}