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

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

AtCoder ABC 292 C - Four Variables (茶色, 300 点)

AtCoder AtCoder300点 ABC-C 茶色diff 数学(整数問題) 全探索 データ構造テク：前処理 バケット 【問題集】整数変数の式で表された条件を扱う探索 入力が定数個 制約:数値が10^6以下 集計処理 半分全列挙 ある値を固定して考える 調和級数

これも最近よく見る「整数の式で表された条件を扱う探索問題」の一味ですね！

問題概要

整数 N が与えられる。

正の整数 A, B, C, D の組であって、

AB + CD = N

を満たすものの個数を求めよ。

制約

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

考えたこと

もし計算時間をまったく気にしなくてよいならば、次のように 4 重の for 文を書くのが楽だろうと思う。

long long res = 0;
for (long long A = 1; A <= N; ++A) {
    for (long long B = 1; B <= N; ++B) {
        for (long long C = 1; C <= N; ++C) {
            for (long long D = 1; D <= N; ++D) {
                 if (A * B + C * D == N) {
                     ++res;
                 }
            }
        }
    }
}

このコードの計算量は O(N^{4}) となる。流石に間に合わない。

　

Step 1： AB と CD に分離する

計算量を落とすにあたって、とても有効な考え方として、「ある量を固定して考える」というのがある。

今回は、 A, B の 2 つを固定して考えてみよう。残りの変数 C, D の組の個数を数えることを考える。ここで、 AB + CD = N という式は、

CD = N - AB

というように変形できることに気づく。ここで、もし次のような配列があったら嬉しいなとなる。

  • num[v] ← 正の整数 x, y の組であって、 xy = v を満たすものの個数

この配列がもし求められていたならば、次のコードによって答えが求められるのだ！

long long res = 0;
for (long long v = 1; v <= N; ++v) {
    // A * B = v となる A, B の個数
    long long ab = num[v];

    // C * D = N - v となる C, D の個数
    long long cd = num[N - v];

    res += ab * cd;
}

めっちゃ簡潔になった。この部分だけなら計算量は O(N) まで落ちた。残りは、配列 num[v] を求める方法を考えよう。

　

Step 2：配列 num を求める

色んな方法が考えられるが、一番簡単なのは次のコードで求める方法だと思われる。

vector<long long> num(N+1, 0);
for (long long x = 1; x <= N; ++x) {
    for (long long y = 1; x * y <= N; ++y) {
        ++num[x * y];
    }
}

つまり、 x, y をさまざまに動かして行ったときの、 x * y の値の分布を集計するのだ。

このコードは一見すると、2 重の for 文なので O(N^{2}) の計算量を要するように見えるかもしれない。

しかし、 x の値が大きくなると、 y の動く範囲は小さくなる。具体的には、 x = 1, 2, 3, \dots, N のとき、 y の動く範囲は \displaystyle N, \frac{N}{2}, \frac{N}{3}, \dots, \frac{N}{N} となる。

この総和は

\displaystyle N(1 + \frac{1}{2} + \frac{1}{3} + \dots + \frac{1}{N})

というように表せる。ここで、ABC でも時々登場する事実なのだが、実は

\displaystyle 1 + \frac{1}{2} + \frac{1}{3} + \dots + \frac{1}{N} = O(\log N)

であることが知られている。よって、

\displaystyle N(1 + \frac{1}{2} + \frac{1}{3} + \dots + \frac{1}{N}) = O(N \log N)

となる。つまり、配列 num を求めるコードの計算量は O(N \log N) なのだ。

なお、上の事実がなぜ成り立つのかについて、関心のある方は次の記事など読んでみるとよいと思う。

manabitimes.jp

まとめ

以上の解法をまとめる。

  1. 配列 num(num[v] は xy = v となる正の整数 x, y の個数を表す) を求める
  2. 配列 num を用いて数え上げる

計算量は O(N \log N) となる。

コード

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

int main() {
    long long N;
    cin >> N;
    
    // 配列 num を求める
    vector<long long> num(N+1, 0);
    for (long long x = 1; x <= N; ++x) {
        for (long long y = 1; x * y <= N; ++y) {
            ++num[x * y];
        }
    }
    
    // 答えを求める
    long long res = 0;
    for (long long v = 1; v <= N; ++v) {
        res += num[v] * num[N - v];
    }
    cout << res << endl;
}