AtCoder Beginner Contest 304 の F 問題に出てきた約数包除のサンプルコード

問題の元ネタはこちら。
ABC の F 問題で、約数系包除原理の問題が出るような時代になったみたいです。

atcoder.jp

約数系包除原理で検索すると、見つかる有名どころの解説記事はここらへんでしょうか。本記事ではあまり包除原理自体の説明には踏み込まないので、仕組みを知りたい方はこちらの記事に飛ぶのがおすすめです。

qiita.com

以下、ABC304F の AC コード を少しだけ書き換えたものを置いておきます。
約数包除系問題の解答コードの雛型として利用することを想定しています。

from itertools import combinations
import math

def dividers_pie(n):
    # 前処理のエラトステネスの篩
    prime_flg = [True] * (n + 1)
    prime_flg[0] = False
    prime_flg[1] = False
    for i in range(2, math.ceil(n ** 0.5) + 1):
        if not prime_flg[i]:
            continue
        for j in range(2, n // i + 1):
            prime_flg[i * j] = False
    primes = [i for i, flg in enumerate(prime_flg) if flg]

    # 素因数と、素因数の個数を求めておく
    prime_factors = [p for p in primes if n % p == 0]
    pf_count = len(prime_factors)

    ans = 0
    # 場合によっては、このループを 1 始まりなどにする必要あり。
    # 元ネタの ABC304F では 1 始まりにしてた
    for i in range(0, pf_count + 1):
        for choosed_pfs in combinations(prime_factors, i):
            factor = 1
            for pf in choosed_pfs:
                factor *= pf
            if factor > n:
                continue
            # ここも問題によりけり。ABC 304F では符号は逆
            sign = 1 if i % 2 == 0 else -1
            # あとは hogehoge を問題に合わせて求める
            ans += sign * hogehoge

    return ans

計算量につきましては、エラトステネスの篩の部分が  O(n \times log(log(n))) ほど。

    for i in range(0, pf_count + 1):
        for choosed_pfs in combinations(prime_factors, i):

ここのループが、  O(2^{x}) (ただし、 x n の素因数の個数)となりますが、  n \leq 10^{5} のとき  x \leq 6 n \leq 10^{6} のとき  x \leq 7 なので大したループにはなりません。