AtCoder ABC 254 D - Together Square: 高速な素因数分解(SPF利用)を用いたO(NloglogN)解法

https://atcoder.jp/contests/abc254/tasks/abc254_d

解説のいくつかで触れられている解法とほぼ同じです。
平方数の性質を考えて、全ての$x$について素因数分解を行い、候補となる$y$を数えます。

単一の数に対する素因数分解は$O(\sqrt(N))$ですが、$N$までの事前計算を$O(Nlog log N)$で行っておくことで$O(log log(N))$で素因数分解が可能になります。

# 平方数の性質
ある数$x$が$k$の平方数とは$x = k^2$です。$k$は$k = p_1^{n_1} \times p_2^{n_2} \times ... \times p_n^{n_a}$というように、$a$個の素因数($p$)で表せるため、$x$は$x = k^2 = k = p_1^{2n_1} \times p_2^{2n_2} \times ... \times p_n^{2n_a}$で表すことができます。すなわち、<b>平方数とは素因数分解した際に、すべての素因数が偶数個含まれている</b>数です。

# 積が平方数になるときと平方数にならないとき
まず、２つの正の整数$x,y$の積を考えます。この時、$a$個の素因数を使って$x \times y = p_1^{n_1 + m_1} \times p_2^{n_2 + m_2} \times ... \times p_a^{n_a + m_a}$と表せます。$n$と$m$は対応する素因数$p_i$を$a$が$n_i$個、$b$が$m_i$個持つことを表します。この時、各$n_i$と$m_i$は0でも良いです。例えば、$12 \times 7 = 2^{2+0} \times 3^{1+0} \times 7^{0+1} $です。

- 平方数 × 平方数 = 平方数
平方数$x, y$同士の積は必ず平方数となります。それぞれに含まれる素因数は偶数個であるため、積に含まれる素因数の数は必ず偶数となるためです。

- 平方数 x 非平方数 = 非平方数
平方数$x$と非平方数$y$の積は必ず非平方数となります。$y$は奇数個しか含まれない素因数が一つ以上含まれます。$x$にはその素因数は$0$あるいは偶数個含まれているため、積に奇数個のある素因数が含まれるため非平方数となります。

- 非平方数 x 非平方数 = ???
非平方数$x,y$の積は両方のケースがあり得ます。例を示します。
1: $12 \times 3$は$2^{2+0} \times 3^{1+1} = 2^{2} \times 3^{2} = 6^2 $で平方数になります。
2: $2 \times 3$は$2^{1+0} \times 3^{0+1} = 2^1 \times 3^1 $で平方数ではありません。

# この問題の考え方
制約: $x,y \leq 2e5$において、すべて数え上げる方針を考えます。

まず、1から$N$に含まれるすべての平方数を$sq$(squares)持っておきます。この上で、いくつか例を見ていきます。

$x=1$のとき、すべての$sq$が候補になります。1はすべての素因数の数に影響を与えないためです。

$x=2$の時を考えます。この時、因数$2^1$を含むため、$y$は因数$2^1$を含まなければ素因数$2$が偶数個にならないため、因数$2$を含み積が$N$を超えない$sq$内の要素との積がすべての候補です。

$x=4$の時を考えます。$4$は平方数のため、$y$は平方数でなければ積は平方数になりません。このため、$4$との積が$N$を超えない$sq$内の要素との積は条件を満たします。

$x=54$の時を考えます。この時、$54 = 2^1 \times 3^3$であるため、$y$は因数として$2^1 \times 3^1$を含まなければ積は素因数$3$と$5$が偶数個にならないため、$3 \times 5$ との積が$N$を超えないような$sq$の要素がすべての候補です。

# 実装方針
- 平方数$sq$の数え上げは$\sqrt(N)$までの平方を管理できます。$\sqrt(N)$つまり今回の制約で含まれる数は$447$個です。
- $x$の素因数の数を求めるのは素因数分解をすれば良いです。
この上で、
- すべての$x$について、素因数分解を行い、奇数個しか含まれない因数の積を求めます。これを$c$とします。素因数分解は工夫をせずに$O(\sqrt(N))$で可能です。$x$すべてについては$O(N\sqrt(N))$です。
- $sq$の各要素と$c$との積が$N$を超えないものを数えます。途中の打ち切りを入れなくても$2e5 \times 5e2 < 1e8$程度で列挙できます。

これらを実装して以下のようにできます。
https://atcoder.jp/contests/abc254/submissions/32292042

# 実装の工夫1: 素因数分解の高速化
素因数分解があるため、$x$の素因数分解に合計で$O(N\sqrt(N))$がかかります。
エラトステネスの篩を工夫したSPF(small prime factor)というテクニックでこれを事前計算$O(NlogN)$, 各クエリ$O(logN)$にします。

https://www.geeksforgeeks.org/prime-factorization-using-sieve-olog-n-multiple-queries/

SPFはエラトステネスの篩と同じ要領で計算を行います。よく、エラトステネスの篩の解説では素数ではない数・つまり合成数に素数ではないチェック(True or False)を付ける説明が用いられますが、この要領で、各合成数について`最小の約数`を保持します(素数は0を持ちます)。この前処理はエラトステネスの篩と同様に$O(N log log N)$で動作します。

クエリを行う際は入力の値をその最小の約数で割り続ければ良く、$O(logN)$で動作します。

※SPFを使った因数分解は入力が比較的小さな値を複数回クエリする場合は有効ですが、単一の大きい値(1e9)などを因数分解する際は事前計算の時間・空間計算量が大きくなるので利用できません

<details><summary>SPF部分の実装(Python)</summary><div>

```Python
class SPF():
    def __init__(self, n):
        self.spf = [0] * (n+1)
        self.isPrime = [True] * (n+1)
        self.isPrime[0] = False
        self.isPrime[1] = False
        self.prime = []
        for i in range(2, n + 1):
            if self.isPrime[i]:
                self.prime.append(i)
                self.spf[i] = i
            for p in self.prime:
                if i * p >  n: break
                if p > self.spf[i]: break
                self.isPrime[i * p] = False
                self.spf[i * p] = p
    def factrication(self, x):
        cnt = defaultdict(int)
        assert  x > 0
        if x == 1: return tuple( [(1,1)] )
        while x != 1:
            cnt[self.spf[x]] += 1
            x //= self.spf[x]
        res = []
        for k in cnt: res.append( (k, cnt[k]) )
        return res
```
</div></details>

# 実装の工夫2: 条件を満たすsqの個数のカウント
$x$の奇数の素因数の積との積が$N$を超えない個数をカウントするのには累積和をとっておき$[i]$を$i$までに含まれる平方数の数としておけば、$O(1)$で個数を求められます。今回の場合は素因数分解のパートの計算量が多いですが、このように工夫できます。


# 最終的なコード
https://atcoder.jp/contests/abc254/submissions/32292393


<details><summary>実装(Python)</summary><div>

```Python
from collections import defaultdict
"""
事前計算 O(NlogN)
クエリ O(logN)
で素因数分解をする
"""
class SPF():
    def __init__(self, n):
        self.spf = [0] * (n+1)
        self.isPrime = [True] * (n+1)
        self.isPrime[0] = False
        self.isPrime[1] = False
        self.prime = []
        for i in range(2, n + 1):
            if self.isPrime[i]:
                self.prime.append(i)
                self.spf[i] = i
            for p in self.prime:
                if i * p >  n: break
                if p > self.spf[i]: break
                self.isPrime[i * p] = False
                self.spf[i * p] = p
    def factrication(self, x):
        cnt = defaultdict(int)
        assert  x > 0
        if x == 1: return tuple( [(1,1)] )
        while x != 1:
            cnt[self.spf[x]] += 1
            x //= self.spf[x]
        res = []
        for k in cnt: res.append( (k, cnt[k]) )
        return res

# nを素因数分解して、[素因数, その個数]のリストを作る
def factorization(n):
    arr = []
    temp = n
    for i in range(2, int(-(-n**0.5//1))+1):
        if temp%i==0:
            cnt=0
            while temp%i==0:
                cnt+=1
                temp //= i
            arr.append([i, cnt])
    if temp!=1: arr.append([temp, 1])
    if arr==[]: arr.append([n, 1])
    return arr
class cumSum1D(object):
    sdat = []
    def init(self):
        pass
    def load(self, l):
        import itertools
        self.sdat = list(itertools.accumulate(itertools.chain([0], l)))
    def query(self, l, r):
        """
        query [l, r)
        """
        # assert l < r
        return self.sdat[r] - self.sdat[l]
n = int(input())
spf = SPF(n)
squares = [] # i**2がn以下の平方数を全列挙
i = 1
cnt = [0] * (n+1)
while i**2 <= n:
    cnt[i**2] += 1
    i += 1
ans = 0
cum = cumSum1D()
cum.load(cnt)
# xを総当たりする
for i in range(1, n+1):
    fact = spf.factrication(i)
    odd = 1
    # 素因数分解を行い、素因数が奇数個含まれている 素因数の積を計算する
    for a, cnt in fact:
        if cnt % 2 == 1: odd *= a
    ans += cum.query(0, n // odd + 1)
print(ans)
```
</div></details>