add multiplicity flag

cmd/primeOmega.go

@@ -23,8 +23,9 @@ import (
 )
 
 var primeOmegaN uint
+var primeOmegaMul bool
 
-func primeOmegaUpdateMultiples(sieve []uint, p uint, n uint) {
+func primeOmegaUpdateMultiples(sieve []uint, p uint, n uint, multiplicity bool) {
 	for q := p; ; q *= p {
 		// omega(a*b) = omega(a) + omega(b) if gcd(a,b) = 1
 		for i := 2 * q; i < n; i += q {
@@ -36,6 +37,10 @@ func primeOmegaUpdateMultiples(sieve []uint, p uint, n uint) {
 		if p*q >= n {
 			break
 		}
+
+		if multiplicity {
+			sieve[p*q] = 1 + sieve[q]
+		}
 	}
 }
 
@@ -54,7 +59,7 @@ func primeOmegaSieve(n uint) chan uint {
 			}
 
 			sieve[i] = 1
-			primeOmegaUpdateMultiples(sieve, i, n)
+			primeOmegaUpdateMultiples(sieve, i, n, primeOmegaMul)
 			ch <- sieve[i]
 		}
 
@@ -103,4 +108,6 @@ func init() {
 
 	primeOmegaCmd.Flags().UintVarP(&primeOmegaN, "limit", "n", 0, "upper limit")
 	primeOmegaCmd.MarkFlagRequired("limit")
+
+	primeOmegaCmd.Flags().BoolVarP(&primeOmegaMul, "with-multiplicity", "m", false, "count prime factors with multiplicity")
 }