diff --git a/README.md b/README.md index 4265ac9..8d8b6b7 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ This creates the `youngkit` executable in the current directory. The command has twelve subcommands. Their usage forms are: ```text -youngkit count N [--max-part K] [--distinct] +youngkit count N [--max-part K] [--largest K] [--distinct] youngkit list N [--max-part K] [--limit L] youngkit rank SHAPE youngkit unrank N RANK @@ -49,7 +49,7 @@ Counts and ranks are exact integers. Shape and tableau commands accept `0` for t ## Commands -`count` counts partitions of a nonnegative integer, optionally limiting the largest part and requiring distinct parts. +`count` counts partitions of a nonnegative integer, optionally limiting the largest part and requiring distinct parts. `--largest K` instead counts the partitions whose largest part is exactly `K`, and an extra `largest` line shows `K`; it cannot be combined with `--max-part`. `list` prints the first partitions in reverse lexicographic order, with an optional largest-part bound and result limit. diff --git a/internal/cli/cli_test.go b/internal/cli/cli_test.go index 0afc6b3..f6ec0d5 100644 --- a/internal/cli/cli_test.go +++ b/internal/cli/cli_test.go @@ -94,7 +94,7 @@ func TestEveryCommandReportsItsExactUsageLine(t *testing.T) { args []string want string }{ - {[]string{"count"}, "usage: youngkit count N [--max-part K] [--distinct]"}, + {[]string{"count"}, "usage: youngkit count N [--max-part K] [--largest K] [--distinct]"}, {[]string{"list"}, "usage: youngkit list N [--max-part K] [--limit L]"}, {[]string{"rank"}, "usage: youngkit rank SHAPE"}, {[]string{"unrank"}, "usage: youngkit unrank N RANK"}, @@ -176,7 +176,7 @@ func TestOptionsAreCheckedBeforeArityAndPositionalParsing(t *testing.T) { "invalid --max-part \"0\" (expected an integer from 1 to 5000)", 0) requireDispatchError(t, []string{"count", "--unknown"}, nil, "unknown option \"--unknown\"", 0) requireDispatchError(t, []string{"count", "bad", "extra"}, nil, - "usage: youngkit count N [--max-part K] [--distinct]", 0) + "usage: youngkit count N [--max-part K] [--largest K] [--distinct]", 0) requireDispatchError(t, []string{"rank", "2,3", "extra"}, nil, "usage: youngkit rank SHAPE", 0) } diff --git a/internal/cli/cmd_count.go b/internal/cli/cmd_count.go index 703fb0c..520e9aa 100644 --- a/internal/cli/cmd_count.go +++ b/internal/cli/cmd_count.go @@ -12,6 +12,7 @@ import ( func RunCount(args []string, io base.IO) (string, error) { specs := []base.Spec{ {Name: "--max-part", Lo: 1, Hi: 5000, Fallback: 0}, + {Name: "--largest", Lo: 1, Hi: 5000, Fallback: 0}, {Name: "--distinct", Flag: true}, } parsed, err := base.Parse(args, specs) @@ -19,13 +20,17 @@ func RunCount(args []string, io base.IO) (string, error) { return "", err } if len(parsed.Args) != 1 { - return "", base.Fail("usage: youngkit count N [--max-part K] [--distinct]") + return "", base.Fail("usage: youngkit count N [--max-part K] [--largest K] [--distinct]") } n, err := base.ParseInt("n", parsed.Args[0], 0, 5000) if err != nil { return "", err } + largest := parsed.Values["--largest"] + if largest != 0 && parsed.Values["--max-part"] != 0 { + return "", base.Fail("--largest and --max-part exclude each other") + } displayedMaxPart := "any" maxPart := parsed.Values["--max-part"] @@ -45,7 +50,11 @@ func RunCount(args []string, io base.IO) (string, error) { base.Line("n", strconv.FormatInt(n, 10)), base.Line("max-part", displayedMaxPart), base.Line("distinct", distinctText), - base.Line("count", total.String()), } + if largest != 0 { + total = count.PartitionsWithLargest(int(n), int(largest), distinct) + lines = append(lines, base.Line("largest", strconv.FormatInt(largest, 10))) + } + lines = append(lines, base.Line("count", total.String())) return strings.Join(lines, "\n") + "\n", nil } diff --git a/internal/cli/count_largest_test.go b/internal/cli/count_largest_test.go new file mode 100644 index 0000000..ed79451 --- /dev/null +++ b/internal/cli/count_largest_test.go @@ -0,0 +1,42 @@ +package cli_test + +import "testing" + +func TestCountLargestMatchesSmallValues(t *testing.T) { + requireCommandSuccess(t, []string{"count", "6", "--largest", "3"}, + "n 6\nmax-part any\ndistinct no\nlargest 3\ncount 3\n") + requireCommandSuccess(t, []string{"count", "6", "--largest=6"}, + "n 6\nmax-part any\ndistinct no\nlargest 6\ncount 1\n") + requireCommandSuccess(t, []string{"count", "6", "--largest", "1"}, + "n 6\nmax-part any\ndistinct no\nlargest 1\ncount 1\n") +} + +func TestCountLargestWithDistinctParts(t *testing.T) { + requireCommandSuccess(t, []string{"count", "10", "--distinct", "--largest", "5"}, + "n 10\nmax-part any\ndistinct yes\nlargest 5\ncount 2\n") + requireCommandSuccess(t, []string{"count", "6", "--largest", "1", "--distinct"}, + "n 6\nmax-part any\ndistinct yes\nlargest 1\ncount 0\n") +} + +func TestCountLargestAboveNOrOfEmptyPartitionIsZero(t *testing.T) { + requireCommandSuccess(t, []string{"count", "4", "--largest", "5"}, + "n 4\nmax-part any\ndistinct no\nlargest 5\ncount 0\n") + requireCommandSuccess(t, []string{"count", "0", "--largest", "1"}, + "n 0\nmax-part any\ndistinct no\nlargest 1\ncount 0\n") +} + +func TestCountLargestLastValueWins(t *testing.T) { + requireCommandSuccess(t, []string{"count", "6", "--largest=bad", "--largest", "2"}, + "n 6\nmax-part any\ndistinct no\nlargest 2\ncount 3\n") +} + +func TestCountLargestErrors(t *testing.T) { + requireDispatchError(t, []string{"count", "6", "--largest", "0"}, nil, + "invalid --largest \"0\" (expected an integer from 1 to 5000)", 0) + requireDispatchError(t, []string{"count", "6", "--largest"}, nil, "--largest needs a value", 0) + requireDispatchError(t, []string{"count", "6", "--largest", "2", "--max-part", "3"}, nil, + "--largest and --max-part exclude each other", 0) + requireDispatchError(t, []string{"count", "6", "--max-part=3", "--largest=2"}, nil, + "--largest and --max-part exclude each other", 0) + requireDispatchError(t, []string{"list", "6", "--largest", "2"}, nil, "unknown option \"--largest\"", 0) +} diff --git a/internal/cli/docs_test.go b/internal/cli/docs_test.go index a6342f6..a111660 100644 --- a/internal/cli/docs_test.go +++ b/internal/cli/docs_test.go @@ -214,7 +214,7 @@ func TestDocumentationUsageLinesMatchNoArgumentDiagnostics(t *testing.T) { line string args []string }{ - {"youngkit count N [--max-part K] [--distinct]", []string{"count"}}, + {"youngkit count N [--max-part K] [--largest K] [--distinct]", []string{"count"}}, {"youngkit list N [--max-part K] [--limit L]", []string{"list"}}, {"youngkit rank SHAPE", []string{"rank"}}, {"youngkit unrank N RANK", []string{"unrank"}}, diff --git a/internal/count/count_test.go b/internal/count/count_test.go index ab19884..55a36d6 100644 --- a/internal/count/count_test.go +++ b/internal/count/count_test.go @@ -370,7 +370,7 @@ func TestRunUnrankChecksNBeforeMalformedRank(t *testing.T) { } func TestAdaptersUsageAndOptionErrors(t *testing.T) { - checkPartitionRun(t, cli.RunCount, nil, "", "usage: youngkit count N [--max-part K] [--distinct]") + checkPartitionRun(t, cli.RunCount, nil, "", "usage: youngkit count N [--max-part K] [--largest K] [--distinct]") checkPartitionRun(t, cli.RunList, nil, "", "usage: youngkit list N [--max-part K] [--limit L]") checkPartitionRun(t, cli.RunRank, nil, "", "usage: youngkit rank SHAPE") checkPartitionRun(t, cli.RunUnrank, nil, "", "usage: youngkit unrank N RANK") diff --git a/internal/count/largest.go b/internal/count/largest.go new file mode 100644 index 0000000..9c50a83 --- /dev/null +++ b/internal/count/largest.go @@ -0,0 +1,13 @@ +package count + +import "math/big" + +// PartitionsWithLargest returns the number of partitions of n whose largest +// part is exactly k. When distinct is true, no part may be repeated. The empty +// partition of 0 has no largest part, so it is never counted. +func PartitionsWithLargest(n, k int, distinct bool) *big.Int { + if n <= 0 || k <= 0 { + return new(big.Int) + } + return new(big.Int).Sub(Partitions(n, k, distinct), Partitions(n, k-1, distinct)) +}