hogecoder

つたじろう(Tsuta_J) 競技プログラミングの記録

Educational Codeforces Round 8 E: Zbazi in Zeydabad

問題概要

原文 → Problem - E - Codeforces

'.' または 'z' のマスのみからなる  N \times M の盤面が与えられる。この盤面内に存在する「Z 型」の正方形領域がいくつ存在するか求めよ。

Z 型の正方形領域の定義

  • 最も上と最も下の行は全て 'z'
  • 反対角成分は全て 'z'
  • それ以外は 'z' でも '.' でもよい

解説

以下の説明では、座標の上下方向を  x と、左右方向を  y として、 x y 列目の座標を  (x, y) と表しています。

まず、前処理をして以下の値を簡単に得られるようにしておきます。

  •  zl_{ij}: 座標  (i, j) の左に 'z' が何連続しているか
  •  zr_{ij}: 座標  (i, j) の右に 'z' が何連続しているか
  •  zld_{ij}: 座標  (i, j) の左下 (反対角成分に対応) に 'z' が何連続しているか

各座標  (i, j) について、その座標を正方形領域の右上としたときにどれくらいの大きさの Z 型が作れるかを考えます。まず、大きさの最大値は  s = \min(zl_{ij}, zld_{ij}) 以下であることがわかります。あとは距離  s-1 以内に存在する左下の候補を高速に求められれば、この問題が解けたことになります。

では、左下の候補はどのように求めればよいでしょうか?反対角線ごとに独立に処理する、つまりある  A という整数を決めた際に、 i + j = A を満たす座標についてまとめて処理することにします。以下で扱う座標  (i, j) は、全て  i + j = A を満たす (同じ反対角線上に属する) ものとします。

反対角線ごとに、右の列から左の列に向けて平面走査することを考えます。各座標  (i, j) について、その座標を左下として利用するときの、正方形領域の右上の座標としてありえる  y 座標最大  j_{\max} を求めてみましょう。これはすでに「座標の右に 'z' が何連続しているか」を求めていますので、 j_{\max} = j + zr_{ij} - 1 と求められます。また、右上の座標としてありえる  y 座標最小  j_{\min} は明らかに  j です。よって、 j \leq y \leq j + zr_{ij} - 1 の範囲で、座標  (i, j) が左下の座標として用いられる可能性があります。

あとはこれを考慮しつつ SegmentTree などによる高速化を考えます。左下の座標として使用できるときはセグ木に要素を追加しておいて、使用できなくなった場合は削除します。追加してから削除するまでの間に、右上の座標の候補が来る場合があるため、その場合は左下としてありえる候補を区間加算で得ます。

言語化が難しすぎるのでソースコード見てください・・・。

ソースコード

// BIT 略

int lft[3010][3010], rgh[3010][3010], dig[3010][3010];

signed main() {
    int H, W; cin >> H >> W;
    vector<string> vs(H, string(W, '.'));

    for(int i=0; i<H; i++) {
        cin >> vs[i];
    }

    for(int i=H-1; i>=0; i--) {
        for(int j=0; j<W; j++) {
            if(vs[i][j] == '.') lft[i][j] = dig[i][j] = 0;
            else {
                if(j == 0) lft[i][j] = 1;
                else lft[i][j] = lft[i][j-1] + 1;

                if(i == H-1 or j == 0) dig[i][j] = 1;
                else dig[i][j] = dig[i+1][j-1] + 1;
            }
        }
        for(int j=W-1; j>=0; j--) {
            if(vs[i][j] == 'z') rgh[i][j] = (j == W-1) ? 1 : rgh[i][j+1] + 1;
        }
    }

    int ans = 0;
    for(int A=0; A<=H+W; A++) {
        BIT<int> bit(H+W);

        vector< tuple<int, int, int> > tups;
        for(int x=0; x<H; x++) {
            int y = A - x;
            if(y < 0 or y >= W) continue;

            tups.emplace_back(-(y + rgh[x][y] - 1), 0, x);
            tups.emplace_back(- y                 , 1, x);
            tups.emplace_back(- y                 , 2, x);
        }

        sort(tups.begin(), tups.end());
        for(auto e : tups) {
            int x = get<2>(e), y = A - x, query = get<1>(e);
            if(query == 0) bit.add(x+1,  1);
            if(query == 1) ans += bit.sum(x+min(dig[x][y], lft[x][y])) - bit.sum(x);
            if(query == 2) bit.add(x+1, -1);
        }
    }
    cout << ans << endl;
    return 0;
}