半分全列挙した！

問題へのリンク

問題概要

正の整数 $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

これをみると、

のときは、個
- $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 個
のときは、個
- $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;
}

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

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

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

問題概要

制約

考えたこと

配列 P を求める

コード