diff --git a/util/primes.go b/util/primes.go new file mode 100644 index 0000000..8d04823 --- /dev/null +++ b/util/primes.go @@ -0,0 +1,34 @@ +package util + +import ( + "math" +) + +// GeneratePrimes returns the n-th prime number +// its param primes []int is intended to contain previously found prime numbers +// to reduce duplicated work, but still be testable +func GeneratePrimes(primes []int, n int) int { + if len(primes) < 2 { + primes = []int{2, 3} + } + if len(primes) >= n { + return primes[n-1] + } + for i := primes[len(primes)-1] + 2; len(primes) <= n; i += 2 { + // check if i is a prime number by checking if it is divisible by any of the previous values of primes + // stop at the square root of i + for _, v := range primes { + // not a prime, stop this loop + if i%v == 0 { + break + } + if math.Sqrt(float64(i)) < float64(v) { + // add to primes + primes = append(primes, i) + break + } + } + } + + return primes[n-1] +} diff --git a/util/primes_test.go b/util/primes_test.go new file mode 100644 index 0000000..ae7f1ab --- /dev/null +++ b/util/primes_test.go @@ -0,0 +1,52 @@ +package util + +import ( + "fmt" + "testing" +) + +func TestGeneratePrimes(t *testing.T) { + tests := []struct { + previousPrimes []int + n, expected int + }{ + {[]int{2, 3}, 1, 2}, + {[]int{2, 3}, 2, 3}, + {[]int{2, 3}, 3, 5}, + {[]int{2, 3}, 4, 7}, + {[]int{2, 3, 5, 7}, 4, 7}, + {[]int{2, 3, 5, 7}, 5, 11}, + {[]int{2, 3, 5, 7}, 6, 13}, + {[]int{2, 3}, 7, 17}, + {[]int{2, 3}, 8, 19}, + } + for _, test := range tests { + t.Run(fmt.Sprintf("%v-th prime number", test.n), + func(t *testing.T) { + ans := GeneratePrimes(test.previousPrimes, test.n) + if ans != test.expected { + t.Errorf("Expected %v-th prime to be %v, got %v", test.n, test.expected, ans) + } + }, + ) + } +} + +// run go test -bench=. from within this util folder +// Benchmark +func benchGenPrimes(n int, b *testing.B) { + for i := 0; i < b.N; i++ { + GeneratePrimes([]int{2, 3}, n) + } +} + +// Benchmark generating different magnitudes of primes... +func BenchmarkGeneratePrimes10(b *testing.B) { benchGenPrimes(10, b) } +func BenchmarkGeneratePrimes100(b *testing.B) { benchGenPrimes(100, b) } +func BenchmarkGeneratePrimes1000(b *testing.B) { benchGenPrimes(1000, b) } +func BenchmarkGeneratePrimes10000(b *testing.B) { benchGenPrimes(10000, b) } +func BenchmarkGeneratePrimes100000(b *testing.B) { benchGenPrimes(100000, b) } +func BenchmarkGeneratePrimes1000000(b *testing.B) { benchGenPrimes(1000000, b) } + +// takes ~2 minutes on my mac +// func BenchmarkGeneratePrimes10000000(b *testing.B) { benchGenPrimes(10000000, b) }