公式

D - Count Subgrid Sum = K 解説 by physics0523


抜き出す長方領域の上辺と下辺を固定した問題を時間計算量 \(O(W)\) で解き、全体で時間計算量 \(O(H^2W)\) で解くことを考えましょう。

上辺と下辺を固定した際の \(j\) 列目の累積和を \(A_j\) とします。また、 \(B_j=A_1+A_2+\dots+A_j\) とします。これらは、適切な予計算の下で時間計算量 \(O(W)\) で求められます。
\(B_r-B_l\)\(l+1\) 列目から \(r\) 列目までの長方領域に書き込まれた整数の和に対応するので、数えるべきは \(B_r-B_l = K\) かつ \(0 \le l < r \le W\) なる \((l,r)\) の組の個数です ( ただし \(B_0=0\) とします )。

\(B_r-B_l \ge K\) なる組の数から \(B_r-B_l > K\) なる組の数を引くことにすると見通しがよいです。
全ての \(j\) について \(A_j \ge 0\) であることから \(B\) は (広義) 単調増加なので、両方の問題について \(l\) が増加するにつれてありうる最小の \(r\) も (広義) 単調増加します。そして、最小の \(r\) 以上 \(W\) 以下の全ての整数が \(r\) として適切であることも分かります。
よって、両方の問題はどちらも尺取り法を利用して解くことができます。尺取り法の \(1\) 回あたりの時間計算量は \(O(W)\) です。

時間計算量は \(O(H^2W)\) です。

実装例 (C++):

#include<bits/stdc++.h>

using namespace std;
using ll=long long;

ll subp(vector<ll> &a,ll k){
  ll n=a.size();
  ll r1=0,r2=0,res=0;
  for(ll l=0;l+1<n;l++){
    r1=max(r1,l+1); r2=max(r2,l+1);
    while(r1<n && a[r1]-a[l]<k){r1++;}
    while(r2<n && a[r2]-a[l]<=k){r2++;}
    res+=(r2-r1);
  }
  return res;
}

int main(){
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  long long h,w,k;
  cin >> h >> w >> k;
  vector<vector<ll>> a(h+1,vector<ll>(w+1,0));
  for(ll i=1;i<=h;i++){
    string s;
    cin >> s;
    for(ll j=1;j<=w;j++){
      if(s[j-1]=='1'){a[i][j]=1;}
    }
  }
  for(ll i=1;i<=h;i++){
    for(ll j=1;j<=w;j++){
      a[i][j]+=a[i-1][j];
    }
  }

  ll res=0;
  for(ll u=1;u<=h;u++){
    for(ll d=u;d<=h;d++){
      vector<ll> b(w+1,0);
      for(ll j=1;j<=w;j++){
        b[j]=b[j-1]+(a[d][j]-a[u-1][j]);
      }
      res+=subp(b,k);
    }
  }
  cout << res << "\n";
  return 0;
}

投稿日時:
最終更新: