AtCoder Beginner Contest 304 の F 問題に出てきた約数包除のサンプルコード
問題の元ネタはこちら。
ABC の F 問題で、約数系包除原理の問題が出るような時代になったみたいです。
約数系包除原理で検索すると、見つかる有名どころの解説記事はここらへんでしょうか。本記事ではあまり包除原理自体の説明には踏み込まないので、仕組みを知りたい方はこちらの記事に飛ぶのがおすすめです。
以下、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
計算量につきましては、エラトステネスの篩の部分が ほど。
for i in range(0, pf_count + 1): for choosed_pfs in combinations(prime_factors, i):
ここのループが、 (ただし、
は
の素因数の個数)となりますが、
のとき
、
のとき
なので大したループにはなりません。