massively improve performance with dynamic programming

cmd/stirling.go

@@ -19,6 +19,7 @@ package cmd
 import (
 	"fmt"
 	"math/big"
+	"strconv"
 
 	"github.com/spf13/cobra"
 	"scm.dairydemon.net/filifa/mathtools/internal/lib"
@@ -31,13 +32,13 @@ var stirlingBottom string
 var stirlingUnsigned bool
 
 func stirling(cmd *cobra.Command, args []string) {
-	n, ok := new(big.Int).SetString(stirlingTop, 10)
+	n, err := strconv.Atoi(stirlingTop)
-	if !ok {
+	if err != nil {
 		cobra.CheckErr("invalid input " + stirlingTop)
 	}
 
-	k, ok := new(big.Int).SetString(stirlingBottom, 10)
+	k, err := strconv.Atoi(stirlingBottom)
-	if !ok {
+	if err != nil {
 		cobra.CheckErr("invalid input " + stirlingBottom)
 	}
 

internal/lib/stirling.go

@@ -20,49 +20,57 @@ import (
 	"math/big"
 )
 
-// TODO: implement both of these with dynamic programming
+// given a slice where vals[i] = Stirling1(i+k-1, k-1) for a given k, nextStirling1 updates the slice so vals[i] = Stirling1(i+k, k) using the property that Stirling1(n, k) = -(n-1)*Stirling1(n-1, k) + Stirling1(n-1, k-1)
-
+func nextStirling1(k int, vals []*big.Int) {
-func Stirling1(n, k *big.Int) *big.Int {
+	for i := 1; i < len(vals); i++ {
-	if n.Cmp(big.NewInt(0)) == 0 && k.Cmp(big.NewInt(0)) == 0 {
+		n := int64(k + i - 1)
-		return big.NewInt(1)
+		v := big.NewInt(-n)
+		v.Mul(v, vals[i-1])
+		vals[i].Add(vals[i], v)
+	}
 }
 
-	if n.Cmp(big.NewInt(0)) == 0 || k.Cmp(big.NewInt(0)) == 0 {
+func Stirling1(n, k int) *big.Int {
+	if k > n {
 		return big.NewInt(0)
 	}
 
-	newN := new(big.Int).Set(n)
+	vals := make([]*big.Int, n-k+1)
-	newN.Sub(newN, big.NewInt(1))
+	for i := range vals {
+		vals[i] = big.NewInt(0)
+	}
+	vals[0] = big.NewInt(1)
 
-	result := Stirling1(newN, k)
+	for i := 1; i <= k; i++ {
-	result.Mul(result, newN)
+		nextStirling1(i, vals)
-	result.Neg(result)
-
-	newK := new(big.Int).Set(k)
-	newK.Sub(newK, big.NewInt(1))
-
-	result.Add(result, Stirling1(newN, newK))
-	return result
 	}
 
-func Stirling2(n, k *big.Int) *big.Int {
+	return vals[n-k]
-	if n.Cmp(k) == 0 {
-		return big.NewInt(1)
 }
 
-	if n.Cmp(big.NewInt(0)) == 0 || k.Cmp(big.NewInt(0)) == 0 {
+// given a slice where vals[i] = Stirling2(i+k-1, k-1) for a given k, nextStirling2 updates the slice so vals[i] = Stirling2(i+k, k) using the property that Stirling2(n, k) = k*Stirling2(n-1, k) + Stirling2(n-1, k-1)
+func nextStirling2(k int64, vals []*big.Int) {
+	for i := 1; i < len(vals); i++ {
+		v := big.NewInt(k)
+		v.Mul(v, vals[i-1])
+		vals[i].Add(vals[i], v)
+	}
+}
+
+func Stirling2(n, k int) *big.Int {
+	if k > n {
 		return big.NewInt(0)
 	}
 
-	newN := new(big.Int).Set(n)
+	vals := make([]*big.Int, n-k+1)
-	newN.Sub(newN, big.NewInt(1))
+	for i := range vals {
-
+		vals[i] = big.NewInt(0)
-	result := Stirling2(newN, k)
+	}
-	result.Mul(result, k)
+	vals[0] = big.NewInt(1)
-
+
-	newK := new(big.Int).Set(k)
+	for i := 1; i <= k; i++ {
-	newK.Sub(newK, big.NewInt(1))
+		nextStirling2(int64(i), vals)
-
+	}
-	result.Add(result, Stirling2(newN, newK))
+
-	return result
+	return vals[n-k]
 }