素直な問題だったけど、重たかった

問題概要

$H × W$ のグリッドが与えられる。各マスには白黒で塗られている。

黒マスの 3 つ組であって、それがマンハッタン距離の意味で正三角形になっているようなものが何個あるか数え上げよ。

制約

$1 \le H, W \le 300$

考えたこと

$H, W$ の制約が 300 か... 3 乗までは OK ってことね。全部の黒マス 3 つ組についてチェックしていたら 6 乗かかるのでダメ。

こういう問題は、まず、どう扱ったらいいかよくわからない条件を、わかりやすい条件に言い換えることを考えるのがよさそう。

で、ごちゃごちゃやっていくうちに

3 辺のうち 1 辺が斜め 45 度な辺である
その辺を 1 辺とする正方形の対辺上に、三角形のもう 1 つの頂点がある

ことが必要十分条件であることがわかる。これは僕は「マンハッタン距離だから 45 度回転」を特に意識しないでこの結論になったけれども、45 度回転してもこんな感じの結論になりそう。

で、「3 辺のうち 1 辺が斜め 45 度な辺である」となると自由度が一気に下がって考えやすくなる。

まず、斜め 45 度な辺を fix することを考える。そうするともう 1 つの点は「斜め 45 度な区間」上にあることになる。その区間に何個の黒マスがあるのかは、斜めに累積和をとれば求められそう。

...ここまでで大方の方針は立った。斜め 45 度な辺の自由度は

黒マスを 1 個選んで (2 乗オーダー)
そこから斜め 45 度方向に進んでいって得られる各黒マスについて (1 乗オーダー)
もう 1 つの点としてふさわしい黒マスが何個あるかは斜め累積和によって求められる

しかし重大な注意点があって、三角形によっては

斜め 45 度な辺が 2 つある

というものがあって、これをダブルカウントしてしまう。そこで

斜め 45 度な辺が 1 個なものは ×2 して加算
斜め 45 度な辺が 2 個なものは ×1 して加算

しておいて、最後に 2 で割ることにした。

あと細かいこととして、斜め 45 度な辺として、右上から左下へと行くものに限定して数えあることにした。反対方向なものについては、盤面を左右反転して改めて数えることに。

#include <iostream>
#include <vector>
#include <cstring>
#include <algorithm>
using namespace std;

long long sum[2100][2100];

int check(int i, int j, const vector<string> &fi) {
    int n = (int)fi.size();
    int m = (int)fi[0].size();
    if (i < 0 || i >= n || j < 0 || j >= m) return 0;
    if (fi[i][j] == '#') return 1;
    else return 0;
}

long long solve(vector<string> fi) {
    // sum[a][b] := マス (i, j) であって i+j = a, i <= b を満たす範囲内にある黒マスの個数 (斜め累積和)
    memset(sum, 0, sizeof(sum));
    int n = (int)fi.size();
    int m = (int)fi[0].size();
    for (int ij = 0; ij < n+m-1; ++ij) {
        sum[ij][0] = 0;
        for (int i = 0; i < (n+m)*3; ++i) {
            int j = ij - i;
            int add = 0;
            if (i >= 0 && i < n && j >= 0 && j < m && fi[i][j] == '#') add = 1;
            sum[ij][i+1] = sum[ij][i] + add;
        }
    }
    
    long long res = 0;
    for (int i = 0; i < n; ++i) {
        for (int j = 0; j < m; ++j) {
            if (fi[i][j] != '#') continue;
            for (int k = 1; k < n+m; ++k) {
                int i2 = i + k;
                int j2 = j - k;
                if (i2 < 0 || i2 >= n || j2 < 0 || j2 >= m) continue;
                if (fi[i2][j2] != '#') continue;
                
                // 対辺が左上方向にあるもの
                {
                    int ij = i+j - k*2;
                    if (ij >= 0) {
                        if (check(i-k, j-k, fi)) ++res; // 直角二等辺三角形
                        if (check(i, j-k*2, fi)) ++res; // 直角二等辺三角形
                        
                        int a = max(i-k+1, 0);
                        int b = max(i, 0);
                        if (a <= b) res += (sum[ij][b] - sum[ij][a]) * 2; // それ以外
                    }
                }
                
                // 対辺が右下方向にあるもの
                {
                    int ij = i+j + k*2;
                    if (ij <= n+m-2) {
                        if (check(i+k, j+k, fi)) ++res; // 直角二等辺三角形
                        if (check(i+k*2, j, fi)) ++res; // 直角二等辺三角形
                        
                        int a = i+k+1;
                        int b = i+k*2;
                        res += (sum[ij][b] - sum[ij][a]) * 2; // それ以外
                    }
                }
            }
        }
    }
    return res;
}

int main() {
    int H, W;
    cin >> H >> W;
    vector<string> input(H);
    for (int i = 0; i < H; ++i) cin >> input[i];
        
    long long res = 0;
    res += solve(input);

    // 盤面を左右反転
    for (int i = 0; i < H; ++i) reverse(input[i].begin(), input[i].end());
    res += solve(input);
        
    cout << res / 2 << endl;
}

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

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

Tenka1 2018 E - Equilateral (橙色, 700 点)

問題概要

制約

考えたこと