CodeChef-LUCKYBAL : Little Elephant and Balance

問題概要

文字列Tを前半と後半(ただし後半は非空)に適当に分け、前半に含まれる4の数と後半に含まれる7の数が一致するときTはバランスがとれているという。長さL(<10^5)の文字列が与えられるので、Sの部分文字列でバランスがとれているものがいくつあるか求める問題(開始位置と長さが異なっていれば重複して数えることもある)。

解法

文字列Tがバランスがとれている条件を考える。4[i,j]を部分文字列T[i,j]([i,j)の半開区間)に含まれる4の個数とし、7[i,j]も同様に定める。このときTがバランスがとれていると、あるk(<|T|=:$)が存在して4[0,k]=7[k,$]が成り立つ。4[0,k]=7[k,$]=($-k)-4[k,$]なので4[0,k]+4[k,$]=($-k)、すなわち4[0,$]=($-k)である。右辺は[1,$]のすべての値をとりうるので、結局4を含む部分文字列はすべてバランスがとれている。
なので、すべての部分文字列の作り方からバランスがとれていない部分文字列=7のみからなる部分文字列の作り方を引けば良い。

acceptされたコード

#include <cstring>
#include <cstdio>
#include <algorithm>
using namespace std;

typedef long long int64;

const int MAX_N = (int)(1e5);

int N;
char buf[MAX_N + 2];
int accum[MAX_N + 1];

void init(){
	scanf(" %s ", buf);
	N = strlen(buf);
	if(buf[N-1] != '4' && buf[N-1] != '7'){
		buf[N-1] = '\0';
		--N;
	}
}

int64 solve(){
	int64 ans = (int64)N * (N + 1) / 2;
	int cnt = 0;
	for(int i=0; i<N; ++i){
		if(buf[i] == '4'){
			ans -= (int64)cnt * (cnt+1) / 2;
			cnt = 0;
		}
		else{
			++cnt;
		}
	}
	ans -= (int64)cnt * (cnt + 1) / 2;
	return ans;
}

int main(){
	int T;
	scanf(" %d ", &T);
	for(int _=0; _<T; ++_){
		init();
		printf("%lld\n", solve());
	}

	return 0;
}