これが ABC の C 問題だったとは...！！！
典型90問の問 4 が結構近いと思った。

問題へのリンク

問題概要

$H \times W$ のグリッド (メモリにおさまらない規模) が与えられる。そのうちの $N$ 個のマスには飴が置いてある。

次の条件を満たすマスの個数を求めよ。

「そのマスと行または列が等しいマス ( $H+W-1$ 個ある) のうち、飴のあるマスの個数がちょうど $K$ 個である」

制約

$1 \le H, W, K, N \le 10^{6}$
$K \le N$

考えたこと

競プロ典型90問の問 4 と同様に、次の値をあらかじめ前処理しておこう。

yoko[i] ← $i$ 行目に飴マスが何個あるか
tate[j] ← $j$ 列目に飴マスが何個あるか

このとき、マス $(i, j)$ と行または列が等しい飴マスの個数は次のように解釈できる。

マス $(i, j)$ が飴マスでない場合：yoko[i] + tate[j]
マス $(i, j)$ が飴マスである場合：yoko[i] + tate[j] - 1

このことを踏まえて、次の手順で求められることがわかる。次の値を求めていくことにしよう。

$A$ ：yoko[i] + tate[j] == K であるようなマス $(i, j)$ の個数
$B$ ：yoko[i] + tate[j] == K であって、 $(i, j)$ 自身が飴マスであるような個数
$C$ ：yoko[i] + tate[j] == K + 1 であって、 $(i, j)$ 自身が飴マスであるような個数

このとき、答えは $A - B + C$ となる。

まず yoko, tate は $O(N)$ の計算量で求められる。 $A$ は各 $i$ 行に対して tate[j] が K - yoko[i] になるような $j$ を数えることで求められる (tate をヒストグラム化することでできる)。 $B, C$ は $N$ 個の飴マスを順に見ることで $O(N)$ でできる。

全体として計算量は $O(H + W + N)$ となる。

コード

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

int main() {
    // 入力
    long long H, W, K, N;
    cin >> H >> W >> K >> N;
    vector<int> X(N), Y(N);
    for (int i = 0; i < N; ++i) {
        cin >> X[i] >> Y[i];
        --X[i], --Y[i];
    }

    // 各行各列の飴マスの個数
    vector<long long> yoko(H, 0);
    vector<long long> tate(W, 0);
    for (int i = 0; i < N; ++i) {
        yoko[X[i]]++;
        tate[Y[i]]++;
    }

    // tate をヒストグラム化
    vector<long long> num(N + 1, 0);
    for (int j = 0; j < W; ++j) num[tate[j]]++;

    // 各値を求める
    long long A = 0, B = 0, C = 0;
    for (int i = 0; i < H; ++i) {
        if (K >= yoko[i]) A += num[K - yoko[i]];
    }
    for (int i = 0; i < N; ++i) {
        long long sum = yoko[X[i]] + tate[Y[i]];
        if (sum == K) ++B;
        else if (sum == K + 1) ++C;
    }
    cout << A - B + C << endl;
}

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

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

AtCoder ABC 023 C - 収集王 (2Q, 試験管青色)

問題概要

制約

考えたこと

コード