filifa
/
mathtools
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

speed up with dynamic programming

This commit is contained in:
filifa 2025-09-16 23:09:22 -04:00
parent ddbb892748
commit 0ae3f4667f
1 changed files with 34 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

60
cmd/partitions.go
View File

@ -19,6 +19,7 @@ package cmd
import ( import (
"fmt" "fmt"
"math/big" "math/big"
"strconv"
"github.com/spf13/cobra" "github.com/spf13/cobra"
) )
@ -26,50 +27,59 @@ import (
var partitionsN string var partitionsN string
var partitionsK string var partitionsK string
// TODO: use dynamic programming // 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 p(n, k *big.Int) *big.Int { func nextPartition(k int, vals []*big.Int) {
if n.Cmp(big.NewInt(0)) == -1 { for i := 1; i < len(vals); i++ {
return big.NewInt(0) if i-k >= 0 {
} else if n.Cmp(big.NewInt(0)) == 0 { vals[i].Add(vals[i], vals[i-k])
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) func p(n, k int) *big.Int {
x := p(newN, k) if n < 0 {
return big.NewInt(0)
}
newK := new(big.Int).Sub(k, big.NewInt(1)) if k > n {
y := p(n, newK) k = n
}
z := new(big.Int).Add(x, y) vals := make([]*big.Int, n+1)
return z 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) { func partitions(cmd *cobra.Command, args []string) {
n, ok := new(big.Int).SetString(partitionsN, 10) n, err := strconv.Atoi(partitionsN)
if !ok { if err != nil {
cobra.CheckErr("invalid input " + partitionsN) cobra.CheckErr("invalid input " + partitionsN)
} }
var k *big.Int var k int
if partitionsK == "" { if partitionsK == "" {
k = new(big.Int).Set(n) k = n
} else { } else {
var ok bool var err error
k, ok = new(big.Int).SetString(partitionsK, 10) k, err = strconv.Atoi(partitionsK)
if !ok { if err != nil {
cobra.CheckErr("invalid input " + partitionsK) cobra.CheckErr("invalid input " + partitionsK)
} }
} }
if n.Cmp(big.NewInt(0)) == -1 { if n < 0 {
cobra.CheckErr("n must be nonnegative") cobra.CheckErr("n must be nonnegative")
} }
if k.Cmp(big.NewInt(0)) == -1 { if k < 0 {
cobra.CheckErr("k must be nonnegative") cobra.CheckErr("k must be nonnegative")
} }