けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 254 D - Together Square (緑色, 400 点)

AtCoder AtCoder400点 ABC-D 緑色diff 数学(整数問題) 【問題集】整数変数の式で表された条件を扱う探索 解空間:O(N^2)通りの選択肢 制約条件:総和<=K ある値を固定して考える 素因数分解 O(√N)まで考えれば十分 入力が定数個 エラトステネスの篩を用いた素因数分解の列挙 エラトステネスの篩 典型要素を詰め合わせた教育的問題

なんかユーザー解説の数がすごいことになっててビビる！

問題概要

N 以下の正の整数 i, j の組であって、 i \times j が平方数であるようなものの個数を求めよ。

制約

  • 1 \le N \le 2 \times 10^{5}

考えたこと

この手の問題では「ある変数を固定して考える」という常套手段がある！！

今回は i を固定して考えよう。つまり私たちは次の問題を解決したいということになる。

整数 i が与えられる。以下の条件を満たす整数 j の個数を求めよ。

  • ij は平方数である
  • 1 \le j \le N である

次に考えるべきことは、「一体どういうときに ij が平方数になるのだろうか」という部分だ。

ij が平方数になる j を求める

まずは、 i を素因数分解してみよう。

i = p_{1}^{e_{1}} p_{2}^{e_{2}} \dots p_{K}^{e_{K}}

と表せるとする。ここで、 e_{i} が奇数であるような素因数 p_{i} たちを改めて q_{1}, q_{2}, \dots, q_{L} とすると、ある整数 A を用いて

i = A^{2} \times q_{1}q_{2}\dots q_{L}

と表せることになる。具体例をあげると、1440 を素因数分解すると

1440 = 2^{5} \times 3^{2} \times 5^{1}

となる。これを上のように表すと

1440 = (2^{2} \times 3)^{2} \times (2 \times 5)

となる。つまり、 A = 12 ということだ。

このとき、 ij = A^{2} \times q_{1}q_{2}\dots q_{L} j が平方数であるということは、 q_{1}q_{2}\dots q_{L} j が平方数であるということと同値になる。

さらに、 q_{1}q_{2}\dots q_{L} j が平方数であるということは、 j を素因数分解したときに、次の条件と同値である！

  • j は q_{1}, q_{2}, \dots, q_{L} を素因数に持ち、しかもそれらの指数は奇数である
  • j のそれら以外の素因数については、指数が偶数である

これをさらに分かりやすく言い換えると、次のようになる。 Q = q_{1}, q_{2}, \dots, q_{L} とおくことにする。

j は平方数 R^{2} を用いて、 j = QR^{2} と表せる

ここまで来ると、実はすでに効率のよい解法が得られている！！　それを次にまとめる。

解法まとめ

解法をまとめると、次のようになる。

i = 1, 2, \dots, N に対して、 ij が平方数となる j ( 1 \le j \le N) の個数を次のように求める。

  • i を素因数分解したときに、指数が奇数であるような素因数の積を Q とする
  • このとき、 ij が平方数であることは、整数 R を用いて j = QR^{2} と表せることと同値である
  • よって、条件を満たす j の個数は「 \frac{N}{Q} 以下の平方数の個数」に一致する

最後に計算量を見積もる。各 i を素因数分解するのに O(N \log\log N) の計算量を要することから、全体の計算量は O(N\sqrt{N}) となる。

なお、高速素因数分解を用いると、 1, 2, \dots, N の素因数分解を O(N \log\log N) の計算量で求めることもできる。その技法を用いると、全体の計算量も O(N \log\log N) となる。

コード

解法 1：愚直に素因数分解する方法

計算量は O(N \sqrt{N}) となる。提出すると 182ms だった。

#include <bits/stdc++.h>
using namespace std;

// 素因数分解
vector<pair<long long, long long> > prime_factorize(long long n) {
    vector<pair<long long, long long> > res;
    for (long long p = 2; p * p <= n; ++p) {
        if (n % p != 0) continue;
        int num = 0;
        while (n % p == 0) { ++num; n /= p; }
        res.push_back(make_pair(p, num));
    }
    if (n != 1) res.push_back(make_pair(n, 1));
    return res;
}

int main() {
    long long N;
    cin >> N;
    
    long long res = 0;
    for (long long i = 1; i <= N; ++i) {
        // i を素因数分解して、指数が奇数である素因数の積を求める
        const auto &pf = prime_factorize(i);
        long long Q = 1;
        for (auto [p, e] : pf) if (e & 1) Q *= p;
        
        // N/Q 以下の平方数の個数を求める
        long long sq = (long long)sqrt(N/Q);
        res += sq;
    }
    cout << res << endl;
}

解法 2：高速素因数分解を用いる方法

エラトステネスの篩を用いる。 O(N \log\log N) となる。提出すると 33ms だった。

#include <bits/stdc++.h>
using namespace std;

// isprime[n] := is n prime?
// min_factor[n] := the min prime-factor of n

struct Eratos {
    vector<int> primes;
    vector<bool> isprime;
    vector<int> min_factor;

    Eratos(int MAX) : primes(),
                      isprime(MAX+1, true),
                      min_factor(MAX+1, -1) {
        isprime[0] = isprime[1] = false;
        min_factor[0] = 0, min_factor[1] = 1;
        for (int i = 2; i <= MAX; ++i) {
            if (!isprime[i]) continue;
            primes.push_back(i);
            min_factor[i] = i;
            for (int j = i*2; j <= MAX; j += i) {
                isprime[j] = false;
                if (min_factor[j] == -1) min_factor[j] = i;
            }
        }
    }

    // prime factorization
    vector<pair<int,int>> prime_factors(int n) {
        vector<pair<int,int> > res;
        while (n != 1) {
            int prime = min_factor[n];
            int exp = 0;
            while (min_factor[n] == prime) {
                ++exp;
                n /= prime;
            }
            res.push_back(make_pair(prime, exp));
        }
        return res;
    }
};


int main() {
    long long N;
    cin >> N;
    
    // エラトステネスの篩を準備
    Eratos er(N + 1);
    
    long long res = 0;
    for (long long i = 1; i <= N; ++i) {
        // i を素因数分解して、指数が奇数である素因数の積を求める
        const auto &pf = er.prime_factors(i);
        long long Q = 1;
        for (auto [p, e] : pf) if (e & 1) Q *= p;
        
        // N/Q 以下の平方数の個数を求める
        long long sq = (long long)sqrt(N/Q);
        res += sq;
    }
    cout << res << endl;
}