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

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

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

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

問題概要

 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;
}