【典型90問】081 - Friendly Group（★5）

問題リンク

解法

体重・身長の取る範囲が小さいということがとても気になる。
これをどうにか使えないかと考えてみると、二次元平面上に(身長, 体重)のようにプロットして、何らかの方法で条件を満たすそれぞれの範囲内にいくつ点が存在するかを高速に計算できれば間に合いそうだと気づく。
(ただ、5000²は間に合うのか？という疑問はあった)
特定の範囲内の和を高速に求める時には、累積和が使える。今回は平面を考えるので、二次元累積和を使う。
累積和を取った後は、$1 ≦ x ≦ 5000, 1 ≦ y ≦ 5000 $の領域の中で作れる、一辺の長さがK+1の任意の正方形領域に含まれる点の数をそれぞれ計算し、最大値を求めるとそれが答えとなる。

---

実装

二次元累積和は、縦(もしくは横)方向に累積和を取った後、横(もしくは縦)方向にさらに累積和を取ることで計算できる。
累積和を元に正方形内に含まれる点の数を計算するところの実装は少しややこしい。
けんちょんさんのQiitaの記事の説明がわかりやすい。

void solve(long long N, long long K, std::vector<long long> A, std::vector<long long> B) {
    vector<vint> XY(5010, vint(5010, 0));
    REP (i, N) XY[A[i]][B[i]]++;

    REP (i, 5000) {
        REP (j, 5000) {
            XY[i+1][j] += XY[i][j];
        }
    }

    REP (i, 5000) {
        REP (j, 5000) {
            XY[i][j+1] += XY[i][j];
        }
    }


    ll mx = 0;
    // 正方形の左上端の座標は、x=j, y=i
    // cur: 正方形内の点の数
    REP (i, 5000 - K + 1) REP (j, 5000 - K + 1) {
        ll cur = XY[i + K][j + K];
        if (i != 0) cur -= XY[i - 1][j + K];
        if (j != 0) cur -= XY[i + K][j - 1];
        if (i != 0 && j != 0) cur += XY[i - 1][j - 1];

        chmax(mx, cur);
    }

    c(mx)
}