speed up with dynamic programming

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 {
-		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)
+// 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])
+		}
+	}
 }
 
-	newN := new(big.Int).Sub(n, k)
-	x := p(newN, k)
+func p(n, k int) *big.Int {
+	if n < 0 {
+		return big.NewInt(0)
+	}
 
-	newK := new(big.Int).Sub(k, big.NewInt(1))
-	y := p(n, newK)
+	if k > n {
+		k = n
+	}
 
-	z := new(big.Int).Add(x, y)
-	return z
+	vals := make([]*big.Int, n+1)
+	for i := range vals {
+		vals[i] = big.NewInt(0)
+	}
+	vals[0] = big.NewInt(1)
+
+	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")
 	}