Codeforces #213, Division 1, A ( Division 2, C ) : Matrix

問題文

http://codeforces.com/contest/365/problem/C

概要

数字からなる文字列 s が与えられる。
行列 b の ( i, j ) 要素を とする。
行列 b 内部の長方形領域であって、要素の和が a となるものの数を求めよ。

解法

長方形領域の総和は、行の和と列の和の積となっています。
従って、長方形を構成する横の辺（＝行の連続する部分列）を決めたとき、列の総和として要請される値が確定します。
総和を全列挙してソートしておくと、二分探索により特定の値の数を対数時間で求めることができます。
（総和の全列挙も前計算していもす法にすると O( N^2 ) です）
全体で O( N^2 log N ) なのでこれで間に合います。

ただし、行の総和 = 0 = a のときは例外処理が必要になります。
このとき、列の選び方は全て条件を満たすので、総数は 1 〜 N の総和となります。

コード

typedef long long LL;
typedef vector<int> VI;

#define REP( i, m, n ) for ( int i = (int)( m ); i < (int)( n ); ++i )
#define ALL( c ) (c).begin(), (c).end()

#define PB( n ) push_back( n )

#define fst first
#define snd second

int main()
{
	cin.tie( 0 );
	ios::sync_with_stdio( false );

	int a;
	cin >> a;

	string s;
	cin >> s;

	const int N = s.size();

	VI row;
	transform( ALL( s ), back_inserter( row ), bind2nd( minus<int>(), '0' ) );

	VI csum( 1, 0 );
	partial_sum( ALL( row ), back_inserter( csum ) );

	VI sums;
	REP( i, 0, N )
	{
		REP( j, i, N )
		{
			sums.PB( csum[ j + 1 ] - csum[i] );
		}
	}
	sort( ALL( sums ) );

	LL res = 0;
	REP( i, 0, N )
	{
		REP( j, i, N )
		{
			const int r = csum[ j + 1 ] - csum[i];

			if ( !r && !a )
			{
				res += N * ( N + 1 ) / 2;
			}

			if ( !r || a % r )
			{
				continue;
			}

			const int c = a / r;
			auto poss = equal_range( ALL( sums ), c );
			res += distance( poss.fst, poss.snd );
		}
	}

	cout << res << endl;			

	return 0;
}