main.go

package main

import (
	"flag"
	"fmt"
	"math/big"
	"time"
)

/* there are four interesting values:
 * - total permutations of N dice
 * - total unique permutations of these specific values
 * - how many fewer permutations there are, NOT including the lowest run
 * - the number of things in that lowest run
 */
type combos struct {
	perm     int64
	unique   int64
	divisor  int64
	firstrun int64
}
type countMap map[int]*big.Int
type dieMap map[int]int
type faceMap map[int][]int
type comboMap map[int]combos

type results struct {
	sides     int
	roll      int
	keep      int
	faces     faceMap
	counts    countMap
	divisors  dieMap
	extraDice dieMap
	sums      dieMap
	totals    countMap
	results   int
}

func newResults(sides int, roll int, keep int) *results {
	r := results{
		sides:   sides,
		keep:    keep,
		roll:    roll,
		faces:   make(faceMap, 0),
		results: 0,
	}
	rolled := make([]int, keep)
	r.sums = make(dieMap, r.keep*r.sides)
	r.totals = make(countMap, r.keep*r.sides)
	r.counts = make(countMap, r.keep*r.sides)
	r.divisors = make(dieMap, r.keep*r.sides)
	r.extraDice = make(dieMap, r.keep*r.sides)
	subtab(&r, 0, rolled)

	maxCombos := make([]*big.Int, r.keep + 1)
	maxCombo := big.NewInt(1)
	moreDice := r.roll - r.keep
	die := big.NewInt(0)
	for i := 0; i < r.keep; i++ {
		tmp := big.NewInt(0)
		tmp.Add(tmp, maxCombo)
		maxCombos[r.keep - i] = tmp
		die.SetInt64(int64(r.roll - i))
		maxCombo.Mul(maxCombo, die)
	}
	lowDie := big.NewInt(0)
	expt := big.NewInt(0)
	combos := big.NewInt(0)
	divisor := big.NewInt(0)
	shadow := big.NewInt(0)
	subShadow := big.NewInt(0)
	permutations := big.NewInt(0)
	pass := big.NewInt(0)
	result := big.NewInt(0)
	one := big.NewInt(1)
	for idx := 0; idx < r.results; idx++ {
		var total *big.Int
		total = r.totals[r.sums[idx]]
		if total == nil {
			var n big.Int
			r.totals[r.sums[idx]] = &n
			total = &n
		}
		extra := r.extraDice[idx]
		permutations.SetInt64(1)
		lowDie.SetInt64(int64(r.faces[idx][0]))
		expt.SetInt64(int64(moreDice + extra))
		divisor.SetInt64(int64(r.divisors[idx]))
		combos.Div(maxCombos[extra], divisor)
		pass.SetInt64(1)
		/* combos is now the number of combos of the dice which aren't
		 * the lowest die. we are now looking at the question of
		 * how many ways there are for best X of Y dice to be Z or
		 * lower, where Z is the lowest die, X is the length of
		 * the run (extra), and Y is the (rolled-kept) dice plus
		 * the length of the run.
		 *
		 * if X is one (and it can't be lower, there's always a lowest
		 * die), the space is the Y-dimensional space of dimension Z-1.
		 * for each additional X, we need to add a one-lower dimensional
		 * space, multiplied by the number of permutations of how
		 * this could happen. so, the second pass would be (Y-1)
		 * dimensional, with X permutations. the third would be (Y-2)
		 * dimensional, with (X*X-1)/2 permutations. The fourth would
		 * be (Y-3) dimensional, with (X*X-1*X-2)/6 permutations.
		 *
		 * so, we start with the total space "shadowed" by this roll,
		 * which is the whole Z^Y space.
		 */
		shadow.Exp(lowDie, expt, nil)
		lowDie.Sub(lowDie, one)
		for extra > 0 {
			subShadow.Exp(lowDie, expt, nil)
			subShadow.Mul(subShadow, permutations)
			shadow.Sub(shadow, subShadow)
			permutations.Mul(permutations, expt)
			permutations.Div(permutations, pass)
			expt.Sub(expt, one)
			pass.Add(pass, one)
			extra--
		}
		result.Mul(combos, shadow)
		total.Add(total, result)
	}
	return &r
}

func expIdx(sides int, rolled []int) int {
	x := rolled[0]
	for i := 1; i < len(rolled); i++ {
		x *= sides
		x += rolled[i]
	}
	return x
}

var factorials []int64 = []int64{1, 2, 6, 24, 120, 720, 5040, 40320, 362880, 3628800}

func fact(n int) int64 {
	// let us be cautious here
	if n < 1 {
		return 1
	}
	if n <= len(factorials) {
		return factorials[n-1]
	}
	for i := len(factorials); i <= n; i++ {
		factorials = append(factorials, factorials[i-1]*int64(i+1))
	}
	return factorials[n-1]
}

func findCombos(dice []int) combos {
	var div1 int64 = 1
	var div2 int64 = 1
	first := 0
	run := 1
	cur := dice[0]
	for i := 1; i < len(dice); i++ {
		if dice[i] == cur {
			run++
		} else {
			if first != 0 {
				div1 *= fact(run)
			} else {
				first = run
			}
			div2 *= fact(run)
			run = 1
			cur = dice[i]
		}
	}
	if first != 0 {
		div1 *= fact(run)
	} else {
		first = run
	}
	div2 *= fact(run)
	perm := fact(len(dice))
	return combos{
		perm:     perm,
		unique:   perm / div2,
		divisor:  div1,
		firstrun: int64(first),
	}
}

func subtab(r *results, die int, rolled []int) {
	start := 1
	if die > 0 {
		start = rolled[die-1]
	}
	for i := start; i <= r.sides; i++ {
		rolled[die] = i
		if die+1 == r.keep {
			r2 := make([]int, r.keep)
			copy(r2, rolled)
			total := 0
			for _, face := range r2 {
				total += int(face)
			}
			r.faces[r.results] = r2
			r.sums[r.results] = total
			r.totals[total] = big.NewInt(0)
			c := findCombos(r2)
			r.counts[r.results] = big.NewInt(c.unique)
			r.extraDice[r.results] = int(c.firstrun)
			r.divisors[r.results] = int(c.divisor)
			// fmt.Printf("[%v]: %d permutations, %d unique, lowest run %d, %d divisor\n", r2, c.perm, c.unique, c.firstrun, c.divisor)
			r.results++
		} else {
			subtab(r, die+1, rolled)
		}
	}
}

func (r results) show(verbose bool, final bool) {
	for t := r.keep; t <= (r.keep * r.sides); t++ {
		fmt.Printf("%d: %d\n", t, r.totals[t])
	}
}

func main() {
	roll := flag.Int("roll", 4, "number of dice to roll")
	keep := flag.Int("keep", 3, "number of dice to keep")
	sides := flag.Int("sides", 6, "sides per die")
	verbose := flag.Bool("verbose", false, "be verbose")
	quiet := flag.Bool("quiet", false, "be quiet")

	flag.Parse()

	start := time.Now()
	r := newResults(*sides, *roll, *keep)
	end := time.Now()
	diff := end.Sub(start)
	elapsed := diff.Seconds()
	fmt.Printf("%dd%dk%d, %.3fms\n", *roll, *sides, *keep, elapsed*1000)
	if !*verbose && !*quiet {
		r.show(*verbose, true)
	}
}