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

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

競プロ典型 90 問 004 - Cross Sum(★2)

グリッドの縦方向の情報と横方向の情報を前処理するのは典型ですね!!

類題とか

drken1215.hatenablog.com

問題概要

 H \times W のグリッドの各マス目に数値が書かれている。

各マス  (i, j) に対して、「そのマスと行または列が等しいマスの数値の総和」を求めて出力せよ。

制約

  •  1 \le H, W \le 2000

考えたこと

 (i, j) に対して、総和をとるべきマスの個数は  H + W - 1 個になります (下図)。なぜなら

  • 同じ行のマスの個数: W
  • 同じ列のマスの個数: H

となるからです。よって  W + H 個だと考えたくなるのですが、マス  (i, j) 自身はどちらのグループにも属するので重複してしまいます。その重複を除いて  H + W - 1 個になります。

よって、それらのマスの値を愚直に総和を取るのでは、計算量は  O(HW(H + W) の計算量となります。なお、 H + W - 1 -1 の部分は計算量オーダーを考えるときは無視できます。

包除原理

このように  HW 個もの値を計算しないといけない場面では、各値を  O(1) で計算できないかを考えてみるのが有効でしょう。たとえば  H = 3,  W = 4 でマス  (1, 2) (0-indexed) を考えるとき、下図のように

(行  1 の総和) + (列  2 の総和) - (マス  (1, 2) の値)

として計算できることがわかります。なおここでマス  (1, 2) の値を引くように、「重複しているところを足したり引いたりする」ことを包除原理と呼ぶことがあります。

前処理

上の考察から、次の  HW 個の値を予め求めておけば、各マスの答えを  O(1) で計算できることがわかります。

  •  0 の総和
  •  1 の総和
  • ...
  •  H-1 の総和
  •  0 の総和
  •  1 の総和
  • ...
  •  W-1 の総和

これらの値をすべて求める作業は  O(HW) で計算できます。このように「ある種の必要な値」を予め求めておくことで、その後の「各マスごとの計算」を高速にできるようにしておくことを、前処理と呼びます。

アルゴリズム全体の計算量を評価すると、

  • 前処理の計算量: O(HW)
  • その後すべてのマスについて答えを計算: O(HW) ( HW 個のマスについてそれぞれ  O(1) で計算できるため)

となりますので、全体として  O(HW) となります。

コード (C++)

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

int main() {
    // 入力
    int H, W;
    cin >> H >> W;
    vector<vector<int>> A(H, vector<int>(W));
    for (int i = 0; i < H; ++i)
        for (int j = 0; j < W; ++j)
            cin >> A[i][j];

    // 前処理
    vector<int> yoko(H, 0);  // i 行目の総和
    vector<int> tate(W, 0);  // j 列目の総和
    for (int i = 0; i < H; ++i) {
        for (int j = 0; j < W; ++j) {
            yoko[i] += A[i][j];
            tate[j] += A[i][j];
        }
    }

    // 各マス
    for (int i = 0; i < H; ++i) {
        for (int j = 0; j < W; ++j) {
            cout << yoko[i] + tate[j] - A[i][j] << " ";
        }
        cout << endl;
    }
}

コード (Python3)

# 入力
H, W = map(int, input().split())
A = [list(map(int, input().split())) for _ in range(H)]

# 前処理
yoko = list(map(sum, A))
tate = list(map(sum, zip(*A)))

# 各マス
for i in range(H):
    print(' '.join(map(lambda j: str(yoko[i] + tate[j] - A[i][j]), range(W)))