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

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

AOJ 2957 MOD Rush (HUPC 2019 day2-F)

これ好き!!

問題概要

長さ  N の整数列  A_{1}, \dots, A_{N} と、長さ  M の整数列  B_{1}, \dots, B_{M} が与えられる。

すべての  i, j に対する  A_{i} %  B_{j} の値の総和を求めよ。

制約

  •  1 \le N, M, A_{i}, B_{j} \le 2 \times 10^{5}

考えたこと

 A_{i} B_{j} の値が  10^{5} 以下であることがポイントに思えてくる。一般に  b に対して

 r =  A_{i} %  b

の総和がどう求められるかを考えてみよう。この値は  A_{i} b で割った商を  q_{i} とすると、

 r = A_{i} - b \times q_{i}

と求められる。よって、次のようになる。

 \sum_{i=1}^{N} (A_{i} %  b)
 = (\sum_{i=1}^{N} A_{i}) - b \times (\sum_{i=1}^{N} q_{i})

いかにも全体として調和級数ネタになりそうな雰囲気になっていた。つまり、 b に対する  \sum_{i=1}^{N} q_{i} O(\frac{\max A}{b}) くらいでできたら嬉しい。次のように実現できる。

  •  0 以上  b 未満であるような  A の個数 × 0
  •  b 以上  2b 未満であるような  A の個数 × 1
  •  2b 以上  3b 未満であるような  A の個数 × 2
  • ...

これらを合算すればよい。以上から全体として、 T = \max(\max A, \max B) として  O(T \log T) の計算量でできる。

コード

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

const int MAX = 510000;
int main() {
    int N, M;
    cin >> N >> M;
    long long S = 0;
    vector<long long> Anum(MAX, 0);
    vector<long long> Bnum(MAX, 0);
    for (int i = 0; i < N; ++i) { int a; cin >> a; Anum[a]++; S += a; }
    for (int i = 0; i < M; ++i) { int a; cin >> a; Bnum[a]++; }
    vector<long long> sum(MAX + 1, 0);
    for (int i = 0; i < MAX; ++i) sum[i+1] = sum[i] + Anum[i];

    long long res = 0;
    for (long long b = 1; b < MAX; ++b) {
        if (Bnum[b] == 0) continue;
        long long qsum = 0;
        for (long long q = 0; b * (q + 1) <= MAX; ++q) {
            qsum += q * (sum[b * (q + 1)] - sum[b * q]);
        }
        res += Bnum[b] * (S - b * qsum);
    }
    cout << res << endl;
}