CODE THANKS FESTIVAL 2015 H: 穴あきケーキ

問題

code-thanks-festival-2015-open.contest.atcoder.jp

解法

2次元の累積和をとる。こうすると長方形の中の累積和がO(1)で得られる。

kenkoooo.hatenablog.com

後はしゃくとりする。

コード

ライブラリはmamekinさんのコピペです。

#include <cstdio>
#include <iostream>
#include <sstream>
#include <fstream>
#include <iomanip>
#include <algorithm>
#include <cmath>
#include <string>
#include <vector>
#include <list>
#include <queue>
#include <stack>
#include <set>
#include <map>
#include <bitset>
#include <numeric>
#include <limits>
#include <climits>
#include <cfloat>
#include <functional>
using namespace std;

template <class T>
class CumulativeSum {
  int ny, nx;
  vector<vector<T>> sum;

 public:
  CumulativeSum(const vector<vector<T>>& a) {
    ny = a.size();
    nx = a[0].size();
    sum.assign(ny + 1, vector<T>(nx + 1, 0));
    for (int i = 0; i < ny; ++i) {
      for (int j = 0; j < nx; ++j) {
        sum[i + 1][j + 1] = a[i][j] + sum[i][j + 1] + sum[i + 1][j] - sum[i][j];
      }
    }
  }

  T getSum(int y1, int x1, int y2, int x2) {
    if (y1 > y2 || x1 > x2) return 0;
    y1 = max(y1, 0);
    x1 = max(x1, 0);
    y2 = min(y2, ny - 1);
    x2 = min(x2, nx - 1);
    return sum[y2 + 1][x2 + 1] - sum[y1][x2 + 1] - sum[y2 + 1][x1] +
           sum[y1][x1];
  }
};

int main() {
  int H, W, K;
  cin >> H >> W >> K;

  vector<vector<int>> cake(H, vector<int>(W));

  vector<CumulativeSum<int>> csRemover;

  vector<vector<vector<int>>> remover(
      10, vector<vector<int>>(H, vector<int>(W, 0)));
  for (int i = 0; i < H; ++i) {
    for (int j = 0; j < W; ++j) {
      char c;
      cin >> c;
      cake[i][j] = c - '0';
      remover[cake[i][j]][i][j]++;
    }
  }
  for (int i = 0; i < 10; ++i) {
    csRemover.push_back(CumulativeSum<int>(remover[i]));
  }
  CumulativeSum<int> csCake(cake);

  long long ans = 0;
  for (int y1 = 0; y1 < H; ++y1) {
    for (int y2 = y1; y2 < H; ++y2) {
      int x2 = 0;
      for (int x1 = 0; x1 < W; ++x1) {
        while (x2 < W && csCake.getSum(y1, x1, y2, x2) <= K) x2++;
        if (x2 == W) break;

        for (int w = 0; w < 10; ++w) {
          int x3 = x2 + w;
          if (x3 >= W) break;

          int full = csCake.getSum(y1, x1, y2, x3);
          int diff = full - K;
          if (diff > 9) break;

          ans += csRemover[diff].getSum(y1 + 1, x1 + 1, y2 - 1, x3 - 1);
        }
      }
    }
  }

  cout << ans << endl;
}