素因数分解が既に分かっている場合の約数列挙 (Python)

9034 ワード
AtCoder Python アルゴリズム 競技プログラミング 数学 Python テキストリンク

はじめに

AtCoder等の競技プログラミング問題を解いているとき、「素因数分解は既に分かっているから、これを組み合わせて約数を生成したい」ということをやりたくなる場合があります。
その方法について、教わった&考えた内容をまとめました。

問題

このような問題を考えます。

問題
$N$ 個の整数からなる数列 $A$ があります。
数列 $A$ の積 $A_1 * A_2 * ... * A_N$ を $X$ とするとき、 $X$ の約数をすべて出力してください。

制約

  • $N <= 2 * 10^5$
  • $A_i <= 2 * 10^5$
  • $X <= 10^{18}$

入力形式

N
A1 ... AN

Sample

Input

3
2 2 3

Output

1
2
3
4
6
12

素朴な解法

素朴には、まず積 $X$ を求めてから、 $2$ ～ $\sqrt X$ の範囲で試し割りすればよいですが、
試し割りの計算量は $O(\sqrt N)$ のため、積 $X$ が巨大な場合に計算量が大きくなってしまいます。

Xの素因数分解を求める

ここで、積 $X$ を直接求めるのではなく、 積 $X$ の素因数分解 を求めることを考えます。
次の図で示すように、自然数の掛け算は素因数のマージ であることを利用します。

つまり、数列 $A$ の各要素ごとに素因数分解をし、その結果の 指数部を足し合わせる ことでマージします。
これは $O(N * \sqrt {maxA})$ で求めることができ、今回の制約下で高速に求まります。

from collections import defaultdict
def prime_factors_from_list(A):
    """
    数列Aの全要素を掛けた数の素因数分解を求めます。
    """
    # 数列Aの各要素を素因数分解しながらマージする。
    # 素因数の辞書を1つ用意し、各要素について素因数を発見するたびに追加していく。
    prime_factors = defaultdict(int)
    for a in A:
        tmp = a
        factor = 2
        while factor ** 2 <= a and tmp != 1:
            while tmp % factor == 0:
                a_is_prime = False
                prime_factors[factor] += 1
                tmp //= factor
            # 2,3,5,7,9...
            factor += (1 if factor == 2 else 2)
        if tmp != 1:
            prime_factors[tmp] += 1
    return prime_factors

print(prime_factors_from_list([12,15]))
""" (出力結果。 2^2 * 3^2 * 5^1 という素因数分解を表します。)
defaultdict(<class 'int'>, {2: 2, 3: 2, 5: 1})
"""

素因数分解から約数を列挙する

このようにして求めた積 $X$ の素因数分解から、約数リストを求めるにはどうすればよいでしょうか。

たとえば、 $12$ の素因数分解は $2^2 * 3^1$ であり、約数は $1,2,3,4,6,12$ の6つです。
これは次のように、 各素因数についての、指数の選び方のパターン と対応付けることができます。

$1 = 2^0 * 3^0$
$2 = 2^1 * 3^0$
$3 = 2^0 * 3^1$
$4 = 2^2 * 3^0$
$6 = 2^1 * 3^1$
$12 = 2^2 * 3^1$

素因数 $2$ の指数は $0,1,2$ の3通り、 $3$ の指数は $0,1$ の2通りなので、
$3*2=6$ 個の約数が得られています。

この考え方は、 既知の約数それぞれに対して、今注目している素因数を掛けたものを求めて約数リストに追加する という作業の繰り返しとして実装できます。

def divisors_from_prime_factors(prime_factors, need_sort=True):
    """
    素因数分解から約数リストを生成します。(1とその数自身を含む)

    Parameters
    ---
    prime_factors
        素因数分解を表現した、タプルのリスト。
        p1^a1 * p2^a2 であれば、 [(p1,a1),(p2,a2)] 。
    need_sort
        Trueの場合、約数リストをソートして返します。
    """
    # 既知の約数リスト
    div = [1]
    for p,a in prime_factors:
        # 既知の約数それぞれに対して、
        # p^1倍, p^2倍, ... p^a倍 したものを計算して約数リストに追加する
        m = len(div)
        for i in range(m):
            for j in range(1, a+1):
                div.append(div[i] * p**j)
    if need_sort:
        div.sort()
    return div

完成

上記の関数を用いて、次のように冒頭の問題が解けました。

n = int(input())
A = list(map(int, input().split())
prime_factors = prime_factors_from_list(A)
divisors = divisors_from_prime_factors(prime_factors.items())
print(*divisors, sep="\n")