diff --git a/cmd/totient.go b/cmd/totient.go index e71cf72..5c4b38c 100644 --- a/cmd/totient.go +++ b/cmd/totient.go @@ -18,28 +18,29 @@ package cmd import ( "fmt" - "strconv" "github.com/spf13/cobra" ) -func totientSieve(n int) chan int { - totients := make([]int, n) +var totientN uint + +func totientSieve(n uint) chan uint { + totients := make([]uint, n) totients[0] = 0 totients[1] = 1 - for i := 2; i < n; i++ { + for i := uint(2); i < n; i++ { totients[i] = i - 1 } - ch := make(chan int) + ch := make(chan uint) go func() { - for i := 0; i < n; i++ { + for i := uint(0); i < n; i++ { ch <- totients[i] if i == 0 || i == 1 || totients[i] != i-1 { continue } - for j := 2 * i; j < n; j += i { + for j := uint(2 * i); j < n; j += i { totients[j] -= totients[j] / i } } @@ -51,12 +52,7 @@ func totientSieve(n int) chan int { } func totient(cmd *cobra.Command, args []string) { - n, err := strconv.Atoi(args[0]) - if err != nil { - cobra.CheckErr(err) - } - - for v := range totientSieve(n) { + for v := range totientSieve(totientN) { fmt.Println(v) } } @@ -81,4 +77,7 @@ func init() { // Cobra supports local flags which will only run when this command // is called directly, e.g.: // totientCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle") + + totientCmd.Flags().UintVarP(&totientN, "limit", "n", 0, "upper limit") + totientCmd.MarkFlagRequired("limit") }