Generate a CSV-ready table of number attributes.

generate_numbers_csv.py

from __future__ import annotations

from math import isqrt

# An example for a typical downstream usage is available at the bottom of the page.

def is_prime(n: int) -> bool:
    if n < 2:
        return False
    if n == 2:
        return True
    if n % 2 == 0:
        return False

    for d in range(3, isqrt(n) + 1, 2):
        if n % d == 0:
            return False

    return True


def divisor_count(n: int) -> int:
    count = 0

    for d in range(1, isqrt(n) + 1):
        if n % d == 0:
            count += 1
            if d != n // d:
                count += 1

    return count


def is_semiprime(n: int) -> bool:
    if n < 4:
        return False

    for d in range(2, isqrt(n) + 1):
        if n % d == 0:
            return is_prime(d) and is_prime(n // d)

    return False


def is_square(n: int) -> bool:
    root = isqrt(n)
    return root * root == n


def is_cube(n: int) -> bool:
    root = round(n ** (1 / 3))
    return root**3 == n


def is_power_of_2(n: int) -> bool:
    return n > 0 and (n & (n - 1)) == 0


def fibonacci_numbers(limit: int) -> set[int]:
    nums = set()
    a, b = 1, 1

    while a <= limit:
        nums.add(a)
        a, b = b, a + b

    return nums


def triangular_numbers(limit: int) -> set[int]:
    nums = set()
    total = 0
    i = 1

    while total < limit:
        total += i
        if total <= limit:
            nums.add(total)
        i += 1

    return nums


def proper_divisor_sum(n: int) -> int:
    if n == 1:
        return 0

    total = 1

    for d in range(2, isqrt(n) + 1):
        if n % d == 0:
            total += d
            pair = n // d
            if pair != d:
                total += pair

    return total


def digit_sum(n: int) -> int:
    return sum(int(ch) for ch in str(n))


def digital_root(n: int) -> int:
    if n == 0:
        return 0
    return 1 + (n - 1) % 9


def flag(value: bool) -> int:
    return 1 if value else 0


def make_row(n: int, fibs: set[int], triangulars: set[int]) -> list:
    prime = is_prime(n)
    composite = n > 1 and not prime
    semiprime = is_semiprime(n)
    square = is_square(n)
    cube = is_cube(n)
    pow2 = is_power_of_2(n)
    fib = n in fibs
    triangular = n in triangulars
    perfect = proper_divisor_sum(n) == n

    feature_count = sum(
        [
            flag(prime),
            flag(composite),
            flag(semiprime),
            flag(square),
            flag(cube),
            flag(pow2),
            flag(fib),
            flag(triangular),
            flag(perfect),
        ]
    )

    return [
        n,
        flag(prime),
        flag(composite),
        flag(semiprime),
        flag(square),
        flag(cube),
        flag(pow2),
        flag(fib),
        flag(triangular),
        flag(perfect),
        divisor_count(n),
        digit_sum(n),
        digital_root(n),
        feature_count,
    ]


def make_csv_table(max_n: int) -> list[list]:
    """
    Generate a CSV-ready table of number attributes.
    
    Args:
        max_n: maximum integer, inclusive. Must be 1 or greater.
    
    Returns:
        Table where the first row is the header.
    """
    if max_n < 1:
        raise ValueError("max_n must be 1 or greater")

    headers = [
        "n",
        "prime",
        "composite",
        "semiprime",
        "square",
        "cube",
        "pow2",
        "fib",
        "triangular",
        "perfect",
        "divisor_count",
        "digit_sum",
        "digital_root",
        "feature_count",
    ]

    fibs = fibonacci_numbers(max_n)
    triangulars = triangular_numbers(max_n)

    table: list[list] = [headers]

    for n in range(1, max_n + 1):
        table.append(make_row(n, fibs, triangulars))

    return table
    
"""
Typical downstream usage

import csv
from generate_numbers_csv import make_csv_table

table = make_csv_table(100)

with open("numbers.csv", "w", newline="") as f:
    writer = csv.writer(f)
    writer.writerows(table)
"""

if __name__ == "__main__":
    # Simple demo: print rows for 1..10.
    table = make_csv_table(10)
    
    for row in table:
        print(row)