prime factorization using trial division

README.md

Prime factorization is useful for reshaping an array into the highest dimension form where non of the dimension has size of 1.

prime_factorize.py

import math
import random
import numpy as np

def prime_factorize(n: int):
    # trial division
    assert n > 0
    factors = []
    if n < 2:
        return [1]
    while n % 2 == 0:
        factors.append(2)
        n //= 2
    for x in range(2, int(math.ceil(math.sqrt(n)))):
        while n % x == 0:
            factors.append(x)
            n //= x
    if n != 1:
        factors.append(n)
    return factors


def test_prime_factorize():
    for x in range(1, 1000):
        factors = prime_factorize(x)
        assert np.prod(factors) == x
    random.seed(44)
    for _ in range(1000):
        x = random.randint(1000, 1000000)
        factors = prime_factorize(x)
        assert np.prod(factors) == x


def test_prime_factorize_reshape():
    for x in range(1, 1000):
        factors = prime_factorize(x)
        arr = np.random.random(x)
        # success in reshape is the test
        arr.reshape(factors)
    for _ in range(1000):
        x = random.randint(1000, 1000000)
        factors = prime_factorize(x)
        arr = np.random.random(x)
        # success in reshape is the test
        arr.reshape(factors)