【AtCoder】ABC 393 E - GCD of Subset | 緑コーダーが解くAtCoder

atcoder.jp

配点: 475 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: 1253

問題概要
制約
考察
実装例
コメント

問題概要

長さ $N$ の数列 $A$ と $N$ 以下の正整数 $K$ が与えられる。 $i=1,2,\cdots,N$ について以下のクエリを処理せよ。

クエリ : $A_i$ を含むように $K$ 個の要素を $A$ から選ぶとき、それらの最大公約数としてあり得る最大値を求めよ。

制約

$1 \leq K \leq N \leq 1.2 \times 10^6$
$1 \leq A_i \leq 10^6$
入力は全て整数。

考察

まずは K 個を適当に選んでみる

まず、一旦 $A_i$ を選ぶことは忘れて、

$A$ の中から適当に選んできた $K$ 個の要素 $a_1, a_2, \cdots, a_k$ の最大公約数 $\mathrm{gcd}(a_1, a_2, \cdots, a_k)$ が $g$ の倍数

になるかを考えよう。 $\mathrm{gcd}(a_1, a_2, \cdots, a_k) = g$ ということを示すには、ユークリッドの互除法などを用いる必要があり、少々面倒...

これは、以下のように言い換えられる。

条件(#) : $a_1, a_2, \cdots, a_k$ がそれぞれ $g$ の倍数であるか?

理由(クリックで展開)

正整数 $A, B, C$ が次のように素因数分解されるとする。

$A = p^{\alpha_1} q^{\beta_1} r^{\gamma_1} \cdots$
$B = p^{\alpha_2} q^{\beta_2} r^{\gamma_2} \cdots$
$C = p^{\alpha_3} q^{\beta_3} r^{\gamma_3} \cdots$

このとき、 $A, B, C$ の最大公約数 $\mathrm{gcd}(A, B, C)$ は、

$\mathrm{gcd}(A, B, C) = p^{\min(\alpha_1, \alpha_2, \alpha_3)} q^{\min(\beta_1, \beta_2, \beta_3)} r^{\min(\gamma_1, \gamma_2, \gamma_3)} \cdots$

と計算できる。これは2つでも、4つ以上でも同様。

ここで、 $\min(\alpha_1, \alpha_2, \alpha_3) \gt 0$ ならば、 $\mathrm{gcd}(A, B, C)$ は少なくとも $p$ の倍数であることが分かる。先ほどの素因数分解の式も合わせると、

$\mathrm{gcd}(A, B, C)$ が $p$ の倍数 $\: \Leftrightarrow \:$ $A, B, C$ がそれぞれ $p$ の倍数

と言い換えることができる。

ここで、以下のような配列を定義してみる。

$\mathrm{multiple}_i :=$ $A$ の要素のうち $i$ の倍数の個数

これが分かっていれば、もしある数 $x$ について $\mathrm{multiple}_x \geq K$ が成り立つならば、最大公約数として $x$ が答えの候補に挙がるということになる。

A_i を含めて考える

では、元の問題に戻って $a_1, a_2, \cdots, a_k$ に $A_i$ を含めることを考えよう。

このとき、条件(#)の成立を仮定するならば、 $A_i$ も $g$ の倍数でなければならない。

言い換えれば「 $g$ は $A_i$ の約数である」ということになるので、素直に考えると「 $A_i$ の約数を列挙して、その数に対する $\mathrm{multiple_i}$ が $K$ 以上であるか」を確認していけばよいということになる(方針☆)。

...しかし、この方針は計算量が厳しい(詳細は後述)。

方針を変えよう。

$\mathrm{multiple_i} \lt K$ となる $i$ は明らかに条件を満たさないので除くとする。

その上で、残った $i$ を小さい順に見ていき、 $0$ で初期化した配列 $\mathrm{max\_g}$ に対して以下のように値を更新していく。

$t$ を正整数として、 $\mathrm{max\_g}_{ti} = i$ とする。

こうすることで、この配列は

$\mathrm{max\_g}_{j} :=$ $A$ の要素のうち倍数として $K$ 個以上含まれる $j$ の約数の中で最大のもの

が格納されることになり、各クエリの答えは $\mathrm{max\_g}_{A_i}$ に一致する。

※ 正直、自分で書いていて日本語がかなり怪しいので、具体例も含めて公式解説配信も参照することをオススメする。

計算量のお話

ここからは計算量について考えたい。以降、 $M = \underset{1 \leq i \leq N}{\max} A_i$ とする。

まず、 $\mathrm{multiple_i}$ の構築について。

これは、予めmapなどに各数の出現回数を記録しておき、 $i = 1, 2, \cdots, M$ に対して(これが今考えている倍数)、 $j$ を $i$ から $M$ まで $i$ ずつ増やしながら $\mathrm{multiple_i}$ に $\mathrm{map}_j$ を加算していくことで実現できる。

この処理は、見かけ上二重forループだが、実は $O(M \log M)$ の計算量で実現可能である。

理由(クリックで展開)

$j$ に関するループの中の処理が呼び出される回数を、全体に対する割合で考えていくと、

$j = 1$ : 毎回呼び出されるので $1$ 。
$j = 2$ : 1回おきに(2回に1回)呼び出されるので $\frac{1}{2}$ 。
$j = 3$ : 2回おきに(3回に1回)呼び出されるので $\frac{1}{3}$ 。

...

$j = k$ : $k$ 回に1回呼び出されるので $\frac{1}{k}$ 。

となり、これを合計すると $\displaystyle \sum_{k = j}^{M} \frac{1}{k}$ と書ける(調和級数の一部)。

ここで、少々唐突だが以下のグラフを考える。

左から $k$ 個目の黄色の長方形は幅 $1$ 、高さ $\frac{1}{k}$ である

ここから、以下のような評価ができる( $M$ 個の長方形と曲線が囲む面積とを比較している)。

$\displaystyle \begin{align} \sum_{k=1}^{M} \frac{1}{k} & \leq \int_{2}^{M+1} \left( \frac{1}{x} + 1 \right) dx + 1 \\ & = \int_{1}^{M} \frac{dx}{x} + 1 \quad (曲線を x 軸方向に -1 平行移動した) \\ & = \log M + 1 \end{align}$

したがって、オーダーでは $\displaystyle \sum_{k = j}^{M} \frac{1}{k} \approx \log M$ となる。

$j$ の外側では $i$ のループが $M$ 回実行されているので、この部分の計算量は $O(M \log M)$ と見積もることができる。

続いて、 $\mathrm{max\_g}$ の構築について。

実はこの部分も先ほどと同じく、「外側で $M$ 回、内側で $\log M$ 回の二重forループ」である(実装例も参考に)。

したがって、この部分も計算量は $O(M \log M)$ なのだ。

入出力にはそれぞれ $O(N)$ 時間かかるので、全体としては $O(N + M \log M)$ の計算量であることが分かった。

制約より、最大でも $N \approx M = 10^6$ 程度なので、一応2秒以内に収まるかなといった感じ。

ただ、結構ギリギリなのでいつものようにcout << endlで答えの改行出力を行うと TLE してしまう。

atcoder.jp

これは、endlによる出力がバッファを吐き出すものである(= 出力ごとに時間がかかる) からである。

nekonenene.hatenablog.com

対策としては、cout << '\n'と書くことによって、バッファに内容を溜め続けて最後に吐き出すようにすることが最も簡単だと思う。

あとは、mapをvectorに置き換えるというのもアリかも。

ちなみに、前述した(方針☆)の計算量についてだが、ある整数 $N$ の約数列挙には一般的に $O(\sqrt{N})$ の計算量がかかる。

algo-logic.info

したがって、 $\mathrm{multiple_i}$ の構築部分に $O(M \sqrt{M})$ かかってしまい、ここが全体のボトルネックとなって TLE するのが濃厚である。

実装例

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

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)
#define repe(i, start, end) for (auto i = (start); (i) <= (end); (i)++)

using ll = long long;

// ======================================== //

const int MAX = 1000100;

int main()
{
    int N, K;
    cin >> N >> K;
    vector<int> A(N);
    unordered_map<int, int> mp;
    rep(i, 0, N) cin >> A[i], mp[A[i]]++;

    vector<int> multiples(MAX, 0);
    repe(i, 1, MAX)
    {
        for (int j = i; j <= MAX; j += i)
        {
            multiples[i] += mp[j];
        }
    }

    vector<int> max_g(MAX, 0);
    repe(i, 1, MAX)
    {
        if (multiples[i] >= K)
        {
            for (int j = i; j <= MAX; j += i)
            {
                max_g[j] = i;
            }
        }
    }

    rep(i, 0, N) cout << max_g[A[i]] << '\n';

    return 0;
}