yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

D - Small GCD / Codeforces Round 911 (Div. 2)

問題

整数  a, b, c に対して、関数  f(a, b, c) を「 a \leq b \leq c となるように並び替えたときの  \gcd (a, b) の値」を返す関数として定義する。

素数 n 個の整数列  a が与えられるので、 1 \leq i \lt j \lt k \leq n であるすべての  i, j, k に対して  f({a}_{i}, {a}_{j}, {a}_{k}) の値の総和を求めよ。

すなわち、  \displaystyle \sum_{i = 1}^{n} \sum_{j = i + 1}^{n} \sum_{k = j + 1}^{n} f({a}_{i}, {a}_{j}, {a}_{k}) の値を求めよ。

入力

まず最初の1行目に、テストケースの個数を表す整数  t が与えられる。

その後、 t 個のテストケースのそれぞれに対して、以下のように入力が与えられる。

  • 最初の1行目には、整数  n が与えられる。
  • 次の1行には、 n 個の整数列  a の要素が順に与えられる。

条件

  •  1 \leq t \leq 10
  •  3 \leq n \leq 8 \cdot {10}^{4}
    • すべてのテストケースにおける  n の値の合計は  8 \cdot {10}^{4} を超えない。
  •  1 \leq {a}_{i} \leq {10}^{5}

出力

各テストケースでの答えを1行ごとに出力すること。

解法

まずはすべての  {a}_{i}, {a}_{j}, {a}_{k} の組み合わせについて  f({a}_{i}, {a}_{j}, {a}_{k}) を求めるという方法が思いつきますが、これは  O({n}^{3}) の計算量がかかり、 n \leq 8 \cdot {10}^{4} であるので実行時間制限に間に合いません。

ここで、  \mathrm{cnt}_{x} を「 {a}_{i} のうち、値が  x の倍数であるものの個数」、 \mathrm{sum}_{x} を「 f({a}_{i}, {a}_{j}, {a}_{k}) の値が  x の倍数となる  (i, j, k) の組み合わせの個数」として定めます。

また、数列  a を昇順に並び替えることとします。

ここで、  \mathrm{cnt}, \mathrm{sum} の情報が  {a}_{1}, {a}_{2}, \cdots, {a}_{i - 1} について入っているものとします。 このとき、  {a}_{i} の値があることにより、  \mathrm{cnt}, \mathrm{sum} がどのように更新されるかということを考えます。

まず、  \mathrm{sum}_{x} は「 f({a}_{i}, {a}_{j}, {a}_{k}) の値が  x の倍数となる  (i, j, k) の個数」と定義しました。 また、  f(p, q, r) p \leq q \leq r のときの  \gcd (p, q) の値として定義されています。

そのため、  q = {a}_{i} としたときに、  f(p, {a}_{i}, r) の値として存在するであろうものとして考えられるのは、  {a}_{i} の約数です。  d {a}_{i} の約数とするとき、この状態で  f(p, {a}_{i}, r) = d となるような  (p, r) の値の組み合わせの個数は

  •  p については  \mathrm{cnt}_{d} で計算されている
  •  r については  a が昇順で並んでいることから  {a}_{i + 1}, \cdots , {a}_{n} (n - i) 個が考えられる

ということから、  \mathrm{cnt}_{d} \cdot (n - i) 個であると計算が行えます。

従って、  \mathrm{sum} の値については  {a}_{i} のすべての約数  d について  \mathrm{sum}_{d} \leftarrow \mathrm{sum}_{d} + \mathrm{cnt}_{d} \cdot (n - i) として更新していくことにより、  {a}_{i} についての情報を足せるということが言えます。

この  \mathrm{sum} の値が更新された状態で、  {a}_{i} のすべての約数  d に対して  \mathrm{cnt}_{d} \leftarrow \mathrm{cnt}_{d} + 1 と更新すれば、 \mathrm{cnt} についても  {a}_{i} の情報を足せると言えます。

以上から、ソートされた数列  a について、  {a}_{1}, {a}_{2}, \cdots , {a}_{n} の順に、  \mathrm{sum}, \mathrm{cnt} のそれぞれを更新していくことにより、数列  a 全体での  \mathrm{sum}, \mathrm{cnt} を算出できます。

さて、ここで、 \mathrm{sum}_{x} が 「  f({a}_{i}, {a}_{j}, {a}_{k}) = x となる  (i, j, k) の個数」であれば、 {a}_{i} \leq {10}^{5} より  x の値も最大で  {10}^{5} であることから求める総和は \begin{align} \displaystyle \sum_{i = 1}^{n} \sum_{j = i + 1}^{n} \sum_{k = j + 1}^{n} f({a}_{i}, {a}_{j}, {a}_{k}) = \sum_{x = 1}^{{10}^{5}} x \cdot \mathrm{sum}_{x} \end{align} として算出することができます。

しかし実際には、 \mathrm{sum}_{x} f({a}_{i}, {a}_{j}, {a}_{k}) の値が  x の倍数となる  (i, j, k) の個数なので、このように単純に計算するのみでは正しい答えとはなりません。

例えば、  f({a}_{i}, {a}_{j}, {a}_{k}) = 6 となる  (i, j, k ) の組が存在する場合、この組の存在によって  \mathrm{sum}_{x} のうち  x = 1, 2, 3, 6 の部分で1組分が加算されることになります。

ここで、  x の値が大きければ大きいほど、「値が  x の倍数となるものの個数」と「値が  x となるものの個数」が等しくなる、ということに注目します。

そのため  \mathrm{sum}_{x} の値を「 f({a}_{i}, {a}_{j}, {a}_{k}) = x となる  (i, j, k) の個数」として扱うために、 x を大きい順に見ていき、答えに  x \cdot \mathrm{sum}_{x} の値を加えたら、 x のすべての約数  d に対して  \mathrm{sum}_{d} \leftarrow \mathrm{sum}_{d} - \mathrm{sum}_{x} と値を更新することを考えます。

この作業を行うことにより、  x を大きい順に見ているため、該当する  x を見る頃には「いま見ている  x は考えられる値の中で最大値」ということになるため、更新されていった  \mathrm{sum}_{x} の値は「値が  x となるものの個数」となります。

従って、この作業を  x = {10}^{5}, {10}^{5} - 1, \cdots , 1 の順に行うことにより、  f({a}_{i}, {a}_{j}, {a}_{k}) の値の総和を求めることができます。

以上をまとめると、

  1. 数列  a を昇順にソートする
  2.  i = 1, 2, \cdots, n に対して以下の作業を繰り返し行う
    •  {a}_{i} のすべての約数  d に対して、  \mathrm{sum}_{d} \leftarrow \mathrm{sum}_{d} + \mathrm{cnt}_{d} \cdot (n - i) と更新する
    •  \mathrm{sum} の更新後、  {a}_{i} のすべての約数  d に対して、  \mathrm{cnt}_{d} \leftarrow \mathrm{cnt}_{d} + 1 と更新する
  3.  x = {10}^{5}, \cdots, 1 に対して以下の作業を繰り返し行う
    • 答えとなる値に  x \cdot \mathrm{sum}_{x} を足す
    •  x のすべての約数  d に対して、  \mathrm{sum}_{d} \leftarrow \mathrm{sum}_{d} - \mathrm{sum}_{x} と更新する

という3つの手順で総和を求めることができます。

整数  x の約数の個数を  d(x) とおくと、これは  O(n \log n + n + \sum d(x)) の計算量で求めることができ、  1 から  {10}^{5} までの整数の約数の個数の総和は  1166750 個なので、これは実行時間制限に間に合うと考えられます。

ただし、整数  x のすべての約数を毎回計算で求めていると、この方法では実行時間制限を超えてしまう可能性があります。 そのため、  1 \leq x \leq {10}^{5} である整数  x のすべての約数を予め求めて保存しておくことを考えます。

ここで、ある整数  q がある素数  p を素因数に持っているとき、 q の約数は「  (q / p) の約数すべてと、 (q / p) の約数すべてに  p を掛けた整数すべて」の中から重複を除いたもの、ということになります。

このことから、C++set 等のデータ構造を用いて、小さい順に整数  x の約数を計算していくことで、  1 \leq x \leq {10}^{5} であるすべての整数の約数を求めて保存することができると言えます。

ソースコード

まず、ある整数の約数をすべて求める prepare_divisors() 関数を以下のように実装しました。

bool visited[int(1e5 + 5)]; // visited はすでにその整数の約数を計算したかどうかを管理する
set<int> divisors[int(1e5 + 5)]; // divisors[i] は整数 i のすべての約数が入る

void prepare_divisors(int num) { // 整数 num の約数を求めて divisors[num] に保存する
  if (visited[num]) { // もし num の約数を既に求めていたら、何もせず return する
    return;
  }

  visited[num] = true; // num の約数を既に求めたことにする
  for (int i = 2; i * i <= num; i++) {
    if (num % i != 0) {  // num の最小の素因数を探す
      continue;
    }

    // i が整数 num の素因数のとき
    prepare_divisors(num / i); // 整数 (num / i) の約数を求める
    auto iter = divisors[num / i].begin(); // 整数 (num / i) の約数を順番に見る
    while (iter != divisors[num / i].end()) {
      divisors[num].insert(*iter); // (num / i) の約数を divisors[num] にも入れる
      divisors[num].insert(*iter * i); // (num / i) の約数を i 倍したものも divisors[num] に入れる
      iter++;
    }
    break; // これ以上 for ループを回す必要がないので break する
  }

  // num が素数の場合も考慮して、 1 と num を改めて divisors[num] に入れる
  divisors[num].insert(1);
  divisors[num].insert(num);
  return;
}

上記の prepare_divisors() により  1 から  {10}^{5} までのすべての整数の約数が取り出せる状態になった上で、求める総和を返す solve() 関数を以下のように実装しました。

set<int> divisors[int(1e5 + 5)];
ll n, a[int(1e5 + 5)], cnt[int(1e5 + 5)], sum[int(1e5 + 5)];

ll solve() {
  // 値の入力
  cin >> n;
  for (int i = 1; i <= n; i++) {
    cin >> a[i];
  }
  sort(a + 1, a + n + 1); // 数列 a をソートする(手順1)

  // a[1] から a[n] の順に、手順2を実行する
  for (ll i = 1; i <= n; i++) {
    auto iter = divisors[a[i]].begin();
    while (iter != divisors[a[i]].end()) {
      sum[*iter] += cnt[*iter] * (n - i); // a[i] の約数に対して、 sum の値を更新する
      cnt[*iter]++; // a[i] の約数に対して、 cnt の値を更新する
      iter++;
    }
  }

  ll ans = 0; // 答えとなる総和

  // 10^5 から 1 の順に、手順3を実行する
  for (ll i = ll(1e5); i >= 1; i--) {
    ans += i * sum[i]; // 答えに i * sum[i] の値を足す
    auto iter = divisors[i].begin();
    while (iter != divisors[i].end()) {
      sum[*iter] -= sum[i]; // i の約数に対して、 sum の値を更新する
      iter++;
    }
  }

  return ans;
}

感想

非常にボリュームのある問題だったと思います。自力で解けて嬉しいです。