TDT4109/Exercise 3/11 - Doble lokker/task11c.py

from math import sqrt

try:
  from common.inputChecking.typeCheck import inputTypeCheck
except ModuleNotFoundError:
  print('Sjekk README.md for hvilke flagg python trenger')
  exit(1)

# O(√n)
def isPrime(n):
  if n < 2: return False
  if n == 2 or n == 3 or n == 5: return True
  limit = int(sqrt(n))
  numberToCheck = 5
  while numberToCheck <= limit:
    if n % numberToCheck == 0: return False
    numberToCheck += 2  # Skip all even numbers
  return True


# Would be O(log₂(n)), but isPrime(n) is used,
# thus it's O(√n)
def findAllFactors(n):
  factors = []
  while not isPrime(n):
    for i in range(2, int(n)):
      if n % i == 0:
        factors.append(i)
        n = n / i
        break
  factors.append(int(n))
  return factors


def factorize(n):
  factors = []
  if isPrime(n):
    factors.append(n)
  else:
    factors = findAllFactors(n)
  return factors


if __name__ == "__main__":
  n = inputTypeCheck('Skriv inn et positivt heltall: ', int)
  factors = factorize(n)
  if len(factors) == 1:
    print(f'{n} er et primtall')
  else:
    print(f'{n} = {" * ".join([str(tall) for tall in factorize(n)])}')