CodeIQ, ディビジョン・サム 復習

CodeIQ の解説のコードを勉強がてらに自分なりに書き直してみる。元のコードはリストを使って、他の言語にもある程度似せて書ける様になっている。自分が書いたコードは、純粋に Python のみを対象にしたもの。

"""Devision Sum.

数学の問題をプログラミングで解こう！「ディビジョン・サム」問題解説
https://codeiq.jp/magazine/2016/12/47476/

###mod 式
https://detail.chiebukuro.yahoo.co.jp/qa/question_detail/q1349018013


###Data Structure.
An dict's key describes prime and dict's value is exponent.

    prime_factors[prime] = exponent
"""
import math


def calc_devision_sum(n, D):
    # 1)
    prime_factors_of_factorial\
        = calc_prime_factors_of_factorial(n)

    # 2)
    prime_factors_of_factorial_to_the_nth_power\
        = calc_prime_factors_of_factorial_to_the_nth_power(prime_factors_of_factorial)

    # 3)
    aliquot_sum_mod_of_factorial_to_the_nth_power\
        = calc_aliquot_sum_mod(prime_factors_of_factorial_to_the_nth_power, D)

    return aliquot_sum_mod_of_factorial_to_the_nth_power


#
# 1)
#
def calc_prime_factors_of_factorial(n):
    """Create an prime factors of factorial."""
    # http://examist.jp/mathematics/integer/kaijyou-soinsu/
    prime_list = sieve_of_eratosthenes(n)
    prime_factors_of_factorial = dict(zip(prime_list, [0] * len(prime_list)))
    for prime in prime_list:
        k = prime
        while k <= n:
            prime_factors_of_factorial[prime] += n // k
            k *= prime
    return prime_factors_of_factorial


def sieve_of_eratosthenes(n):
    """Get prime list less than or equal to n.

    1) エラトステネスのふるいとその計算量
    http://mathtrain.jp/eratosthenes
    "なお，素数定理とかを使えば
    愚直な方法に対して
    もう少し計算量を厳しく評価できますが，
    とりあえずエラトステネスのふるいの方が
    速いことを理解しておいてください。"

    ### Data Structure
    is_prime[index]
      list of boolean,
      an each element describes its index number is prime or not.
    ex)
    is_prime[1] = False
    is_prime[4] = False
    is_prime[7] = True
    """
    is_prime = [True] * (n + 1)
    is_prime[0] = False
    is_prime[1] = False
    # faster than the below code
    # for i in range(2, n) + 1):
    for i in range(2, int(math.sqrt(n)) + 1):
        for j in range(2 * i, n + 1, i):
            is_prime[j] = False
    prime_list = [i for i in range(n + 1) if is_prime[i]]
    return prime_list


#
# 2)
#
def calc_prime_factors_of_factorial_to_the_nth_power(prime_factors_of_factorial):
    return {base: exponent * n for base, exponent in prime_factors_of_factorial.items()}


#
# 3)
#
def calc_aliquot_sum_mod(prime_factors, D):
    """約数の和は等比数列の和で表現できる."""
    aliquot_sum_mod = 1
    for base, exponent in prime_factors.items():
        # 等比数列の和の剰余
        t = ((pow_mod(base, exponent + 1, D) - 1) * inv_mod(base - 1, D)) % D
        # 約数の和の剰余
        aliquot_sum_mod = (aliquot_sum_mod * t) % D
    return aliquot_sum_mod


def pow_mod(a, n, d):
    """Modular exponentiation, a^n mod d.

    "結局、分配則が成り立つので、途中で剰余を撮ってしまえば良い。"
    https://ja.wikipedia.org/wiki/%E5%86%AA%E5%89%B0%E4%BD%99
    """
    # 実は組み込み関数がある。
    # return pow(a, n, d)
    if n == 0:
        return 1
    elif n % 2 == 0:
        return pow_mod((a * a) % d, n // 2, d)
    else:
        return a * pow_mod(a, n - 1, d) % d


def inv_mod(a, d):
    """単位元と逆元.

    # d を法とする a の逆元 (a*x = 1 (mod d) となる x) を求める.
    http://www.geocities.jp/k27c8_math/math/group/identity_inverse.htm
    """
    (p, q) = gcd_ext(a, d)
    return p % d


def gcd_ext(a, b):
    """拡張ユークリッド互除法.

    # 拡張ユークリッドの互除法
    # a*x + b*y = gcd(a,b) となる (x,y) を求める.
    http://www.tbasic.org/reference/old/ExEuclid.html
    　逆元の求め方も書いてある。
    """
    x, y, lastx, lasty = 0, 1, 1, 0
    while b != 0:
        q = a // b
        a, b = b, a % b
        x, y, lastx, lasty = lastx - q * x, lasty - q * y, x, y
    return (lastx, lasty)


#
# main
#
if __name__ == '__main__':
    n = int(input())
    D = 1000003
    print(calc_devision_sum(n, D))