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

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

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

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

drken1215.hatenablog.com

問題概要

 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

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


  •  Ayoko[i] + tate[j] == K であるようなマス  (i, j) の個数
  •  Byoko[i] + tate[j] == K であって、 (i, j) 自身が飴マスであるような個数
  •  Cyoko[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;
}