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

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

AtCoder ARC 107 B - Quadruple (茶色, 400 点)

半分全列挙した!

問題概要

正の整数  N と整数  K が与えられる。以下の条件を満たす正の整数  a, b, c, d の組の個数を求めよ。

  •  1 \le a, b, c, d \le N
  •  a + b - c - d = K

制約

  •  1 \le N \le 10^{5}
  •  -2(N-1) \le K \le 2(N-1)

考えたこと

愚直な方法としては、次のように 4 重ループをする解法が考えられるかもしれない。しかしこれでは  O(N^{4}) の計算量を要する。

long long res = 0;
for (int a = 1; a <= N; ++a) {
    for (int b = 1; b <= N; ++b) {
        for (int c = 1; c <= N; ++c) {
            for (int d = 1; d <= N; ++d) {
                 if (a + b - c - d == K) ++res;
            }
        }
    }
}

よって、工夫が必要になる。このような場面でしばしば有効なのが半分全列挙!!!具体的には、

  •  P \lbrack v \rbrack :=  (a + b = v) となるような  (a, b) の個数
  •  Q \lbrack v \rbrack :=  (c + d = v) となるような  (c, d) の個数

として、次のようにすれば OK。なお、 K -K とで答えは変わらないので、 K \lt 0 の場合は  K = -K としている ( K \ge 0 を満たすものとしている)。

long long res = 0;
for (int v = K; v <= N*2; ++v) {
    res += P[v] * Q[v-K];
}

よって、配列  P, Q さえ求められれば、 O(N) で答えが出せるのだ!!!

なお、配列  P Q はまったく同じものなので、配列  P のみ求める。

配列 P を求める

残る問題は

  •  P \lbrack v \rbrack :=  (a + b = v) となるような  (a, b) の個数

を求めることである。試しに  N = 5 の場合を表にすると、次のようになる。

a \ b 1 2 3 4 5
1 2 3 4 5 6
2 3 4 5 6 7
3 4 5 6 7 8
4 5 6 7 8 9
5 6 7 8 9 10

これをみると、

  •  v = 2, 3, \dots, N+1 のときは、 v-1
    •  a+b = 2 となる  (a, b) は 1 個
    •  a+b = 3 となる  (a, b) は 2 個
    •  a+b = 4 となる  (a, b) は 3 個
    •  a+b = 5 となる  (a, b) は 4 個
    •  a+b = 6 となる  (a, b) は 5 個
  •  v = N+2, \dots, 2N のときは、 2N+1-v
    •  a+b = 7 となる  (a, b) は 4 個
    •  a+b = 8 となる  (a, b) は 3 個
    •  a+b = 9 となる  (a, b) は 2 個
    •  a+b = 10 となる  (a, b) は 1 個

ということがわかる。より簡潔に書くならば、

 P \lbrack v \rbrack = \min(v - 1, 2N+1-v)

と書ける。よって配列  P を求めるのも  O(N) の計算量でできる。

コード

全体の計算量も  O(N) となる。

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

int main() {
    long long N, K;
    cin >> N >> K;
    if (K < 0) K = -K;
    vector<long long> num(N*2+1, 0);
    for (long long n = 2; n <= N*2; ++n) num[n] = min(n-1, N*2+1-n);
    long long res = 0;
    for (long long n = K; n <= N*2; ++n) res += num[n] * num[n-K];
    cout << res << endl;
}