Add count --largest K

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.
 

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)
 }
 

internal/cli/cmd_count.go

@@ -12,20 +12,25 @@ 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)
 	if err != nil {
 		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
 }

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)
+}

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"}},

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")

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))
+}