feat(tricks): prime numbers functions

Signed-off-by: Pouyan Heyratpour <[email protected]>
janstoon · Dec 4, 2024 · cd1859f · cd1859f
1 parent 5b28775
commit cd1859f
Show file tree

Hide file tree

Showing 2 changed files with 163 additions and 0 deletions.
diff --git a/tricks/mathx/arithmetic.go b/tricks/mathx/arithmetic.go
@@ -1,5 +1,7 @@
 package mathx
 
+import "math"
+
 // Gcd calculates the Greatest common divisor of multiple integers
 func Gcd(a int, bb ...int) int {
 	if len(bb) == 0 {
@@ -21,3 +23,57 @@ func Lcm(a int, bb ...int) int {
 
 	return Lcm(a*bb[0]/Gcd(a, bb[0]), bb[1:]...)
 }
+
+// IsPrime checks if the given number (a) is a prime number
+func IsPrime(n int) bool {
+	for i := 2; i <= int(math.Sqrt(float64(n))); i++ {
+		if n%i == 0 {
+			return false
+		}
+	}
+
+	return n > 1
+}
+
+// PrimeFactors find prime factors of the given number (n)
+func PrimeFactors(n int) []int {
+	ff := make([]int, 0, 10)
+
+	divisor := 1
+	for {
+		if n <= 1 {
+			break
+		}
+
+		divisor = nextPrime(divisor)
+		if n%divisor == 0 {
+			ff = append(ff, divisor)
+
+			for n%divisor == 0 {
+				n /= divisor
+			}
+		}
+	}
+
+	return ff[:len(ff):len(ff)]
+}
+
+// NextPrime returns next prime number greater than the given number (n)
+func NextPrime(n int) int {
+	if n < 1 {
+		n = 1
+	}
+
+	return nextPrime(n)
+}
+
+// nextPrime returns next prime number greater than the given number (n)
+func nextPrime(n int) int {
+	for {
+		n++
+
+		if IsPrime(n) {
+			return n
+		}
+	}
+}
diff --git a/tricks/mathx/arithmetic_test.go b/tricks/mathx/arithmetic_test.go
@@ -28,6 +28,7 @@ func TestGcd(t *testing.T) {
 	assert.Equal(t, 1, mathx.Gcd(11, 13))
 	assert.Equal(t, 3, mathx.Gcd(3, 21))
 	assert.Equal(t, 14, mathx.Gcd(42, 28))
+	assert.Equal(t, 1, mathx.Gcd(999999000001, 2))
 
 	assert.Equal(t, 1, mathx.Gcd(1, 1, 1))
 	assert.Equal(t, 5, mathx.Gcd(5, 5, 5))
@@ -60,6 +61,7 @@ func TestLcm(t *testing.T) {
 	assert.Equal(t, 143, mathx.Lcm(11, 13))
 	assert.Equal(t, 21, mathx.Lcm(3, 21))
 	assert.Equal(t, 84, mathx.Lcm(42, 28))
+	assert.Equal(t, 1999998000002, mathx.Lcm(999999000001, 2))
 
 	assert.Equal(t, 1, mathx.Lcm(1, 1, 1))
 	assert.Equal(t, 5, mathx.Lcm(5, 5, 5))
@@ -71,3 +73,108 @@ func TestLcm(t *testing.T) {
 	assert.Equal(t, 126, mathx.Lcm(6, 9, 21))
 	assert.Equal(t, 84, mathx.Lcm(42, 28, 21))
 }
+
+func TestIsPrime(t *testing.T) {
+	assert.True(t, mathx.IsPrime(2))
+	assert.True(t, mathx.IsPrime(3))
+	assert.True(t, mathx.IsPrime(5))
+	assert.True(t, mathx.IsPrime(7))
+	assert.True(t, mathx.IsPrime(11))
+	assert.True(t, mathx.IsPrime(13))
+	assert.True(t, mathx.IsPrime(97))
+	assert.True(t, mathx.IsPrime(127))
+	assert.True(t, mathx.IsPrime(281))
+	assert.True(t, mathx.IsPrime(389))
+	assert.True(t, mathx.IsPrime(499))
+	assert.True(t, mathx.IsPrime(8191))
+	assert.True(t, mathx.IsPrime(524287))
+	assert.True(t, mathx.IsPrime(6700417))
+	assert.True(t, mathx.IsPrime(999999000001))
+
+	assert.False(t, mathx.IsPrime(1))
+	assert.False(t, mathx.IsPrime(4))
+	assert.False(t, mathx.IsPrime(6))
+	assert.False(t, mathx.IsPrime(8))
+	assert.False(t, mathx.IsPrime(9))
+	assert.False(t, mathx.IsPrime(10))
+	assert.False(t, mathx.IsPrime(12))
+	assert.False(t, mathx.IsPrime(68))
+	assert.False(t, mathx.IsPrime(72))
+	assert.False(t, mathx.IsPrime(90))
+	assert.False(t, mathx.IsPrime(93))
+	assert.False(t, mathx.IsPrime(100))
+	assert.False(t, mathx.IsPrime(121))
+	assert.False(t, mathx.IsPrime(6700415))
+	assert.False(t, mathx.IsPrime(6700419))
+	assert.False(t, mathx.IsPrime(999999000000))
+	assert.False(t, mathx.IsPrime(999999000002))
+	assert.False(t, mathx.IsPrime(999999000004))
+
+	assert.False(t, mathx.IsPrime(0))
+	assert.False(t, mathx.IsPrime(-1))
+	assert.False(t, mathx.IsPrime(-2))
+	assert.False(t, mathx.IsPrime(-3))
+	assert.False(t, mathx.IsPrime(-4))
+	assert.False(t, mathx.IsPrime(-5))
+	assert.False(t, mathx.IsPrime(-6))
+	assert.False(t, mathx.IsPrime(-7))
+}
+
+func TestPrimes(t *testing.T) {
+	assert.Equal(t, []int{2, 3}, mathx.PrimeFactors(6))
+	assert.Equal(t, []int{2, 5}, mathx.PrimeFactors(10))
+	assert.Equal(t, []int{2, 3, 5}, mathx.PrimeFactors(30))
+	assert.Equal(t, []int{3, 7}, mathx.PrimeFactors(21))
+	assert.Equal(t, []int{2, 3, 5, 13}, mathx.PrimeFactors(390))
+	assert.Equal(t, []int{2, 3, 5, 7, 13}, mathx.PrimeFactors(2730))
+	assert.Equal(t, []int{2, 3, 5, 7, 13}, mathx.PrimeFactors(40950))
+	assert.Equal(t, []int{2, 3, 5, 7, 13, 17}, mathx.PrimeFactors(46410))
+
+	assert.Equal(t, []int{2}, mathx.PrimeFactors(4))
+	assert.Equal(t, []int{2}, mathx.PrimeFactors(8))
+	assert.Equal(t, []int{2}, mathx.PrimeFactors(16))
+	assert.Equal(t, []int{2}, mathx.PrimeFactors(32))
+	assert.Equal(t, []int{2}, mathx.PrimeFactors(64))
+	assert.Equal(t, []int{2}, mathx.PrimeFactors(128))
+	assert.Equal(t, []int{2}, mathx.PrimeFactors(256))
+	assert.Equal(t, []int{2}, mathx.PrimeFactors(512))
+	assert.Equal(t, []int{2}, mathx.PrimeFactors(1024))
+	assert.Equal(t, []int{2}, mathx.PrimeFactors(1<<15))
+
+	assert.Equal(t, []int{}, mathx.PrimeFactors(1))
+	assert.Equal(t, []int{2}, mathx.PrimeFactors(2))
+	assert.Equal(t, []int{3}, mathx.PrimeFactors(3))
+	assert.Equal(t, []int{5}, mathx.PrimeFactors(5))
+	assert.Equal(t, []int{7}, mathx.PrimeFactors(7))
+	assert.Equal(t, []int{11}, mathx.PrimeFactors(11))
+	assert.Equal(t, []int{13}, mathx.PrimeFactors(13))
+
+	assert.Equal(t, []int{}, mathx.PrimeFactors(0))
+	assert.Equal(t, []int{}, mathx.PrimeFactors(-1))
+	assert.Equal(t, []int{}, mathx.PrimeFactors(-2))
+	assert.Equal(t, []int{}, mathx.PrimeFactors(-3))
+	assert.Equal(t, []int{}, mathx.PrimeFactors(-4))
+	assert.Equal(t, []int{}, mathx.PrimeFactors(-5))
+	assert.Equal(t, []int{}, mathx.PrimeFactors(-6))
+	assert.Equal(t, []int{}, mathx.PrimeFactors(-7))
+	assert.Equal(t, []int{}, mathx.PrimeFactors(-100))
+	assert.Equal(t, []int{}, mathx.PrimeFactors(-121))
+	assert.Equal(t, []int{}, mathx.PrimeFactors(-127))
+}
+
+func TestNextPrime(t *testing.T) {
+	assert.Equal(t, 2, mathx.NextPrime(0))
+	assert.Equal(t, 2, mathx.NextPrime(1))
+	assert.Equal(t, 3, mathx.NextPrime(2))
+	assert.Equal(t, 5, mathx.NextPrime(3))
+	assert.Equal(t, 5, mathx.NextPrime(4))
+	assert.Equal(t, 7, mathx.NextPrime(5))
+	assert.Equal(t, 999999000001, mathx.NextPrime(999999000000))
+
+	assert.Equal(t, 2, mathx.NextPrime(-1))
+	assert.Equal(t, 2, mathx.NextPrime(-2))
+	assert.Equal(t, 2, mathx.NextPrime(-3))
+	assert.Equal(t, 2, mathx.NextPrime(-4))
+	assert.Equal(t, 2, mathx.NextPrime(-5))
+	assert.Equal(t, 2, mathx.NextPrime(-100))
+}