diff --git a/cmd/partitions.go b/cmd/partitions.go index eb5db5c..a45f6fe 100644 --- a/cmd/partitions.go +++ b/cmd/partitions.go @@ -19,6 +19,7 @@ package cmd import ( "fmt" "math/big" + "strconv" "github.com/spf13/cobra" ) @@ -26,50 +27,59 @@ import ( var partitionsN string var partitionsK string -// TODO: use dynamic programming -func p(n, k *big.Int) *big.Int { - if n.Cmp(big.NewInt(0)) == -1 { +// given a slice where vals[i] = p_{k-1}(i) for a given k, nextPartition updates the slice so vals[i] = p_k(i) using the property that p_k(n) = p_k(n-k) + p_{k-1}(n) +func nextPartition(k int, vals []*big.Int) { + for i := 1; i < len(vals); i++ { + if i-k >= 0 { + vals[i].Add(vals[i], vals[i-k]) + } + } +} + +func p(n, k int) *big.Int { + if n < 0 { return big.NewInt(0) - } else if n.Cmp(big.NewInt(0)) == 0 { - return big.NewInt(1) - } else if k.Cmp(big.NewInt(0)) == 0 { - return big.NewInt(0) - } else if k.Cmp(n) == 1 { - return p(n, n) } - newN := new(big.Int).Sub(n, k) - x := p(newN, k) + if k > n { + k = n + } - newK := new(big.Int).Sub(k, big.NewInt(1)) - y := p(n, newK) + vals := make([]*big.Int, n+1) + for i := range vals { + vals[i] = big.NewInt(0) + } + vals[0] = big.NewInt(1) - z := new(big.Int).Add(x, y) - return z + for i := 1; i <= k; i++ { + nextPartition(i, vals) + } + + return vals[n] } func partitions(cmd *cobra.Command, args []string) { - n, ok := new(big.Int).SetString(partitionsN, 10) - if !ok { + n, err := strconv.Atoi(partitionsN) + if err != nil { cobra.CheckErr("invalid input " + partitionsN) } - var k *big.Int + var k int if partitionsK == "" { - k = new(big.Int).Set(n) + k = n } else { - var ok bool - k, ok = new(big.Int).SetString(partitionsK, 10) - if !ok { + var err error + k, err = strconv.Atoi(partitionsK) + if err != nil { cobra.CheckErr("invalid input " + partitionsK) } } - if n.Cmp(big.NewInt(0)) == -1 { + if n < 0 { cobra.CheckErr("n must be nonnegative") } - if k.Cmp(big.NewInt(0)) == -1 { + if k < 0 { cobra.CheckErr("k must be nonnegative") }