2015 TCO 1D BalancedSubstrings

問題

問題文

0, 1 からなる文字列が次の条件を満たすとき，バランスしているという．

文字列全体を棒，1 を重さ $1$ のおもりだと思い，ある文字を持って釣ったときに水平に釣り合う．

文字列 $s$ のバランスしている部分文字列の数を求めよ． (詳しい定義は問題文の英語を読んでください．)

方針

棒の左端を持ったときの $i$ 文字目まで全体のトルク $a_i$ と重さ $b_i$ を事前に計算しておく． すると，棒の $[l,r)$ の領域のトルクは $a_r - a_l$ ，重さは $b_r - b_l$ と $O(1)$ で求まる． 部分文字列がある点を持って水平に釣り合うということは，左右のトルクがつりあう点がそこにあるということ． 持てる場所は整数座標のみなので，$b_r - b_l = 0$ または重心の座標 $\frac{a_r - a_l}{b_r - b_l}$ が整数になるか確かめれば良い． 累積和の計算に $O(n)$，部分文字列の全探索に $O(n^2)$ なので間にあう．

実装

class BalancedSubstrings {
public:
    int countSubstrings(string s){
        int n = s.size();
        int a[3000], b[3000];
        rep(i,n+1) a[i] = b[i] = 0;
        rep(i,n){
            a[i+1] = a[i];
            b[i+1] = b[i];
            if(s[i]=='1'){
                a[i+1] += i;
                b[i+1] += 1;
            }
        }
        int ans = 0;
        rep(r,n+1)rep(l,r){
            int A = a[r]-a[l];
            int B = b[r]-b[l];
            if(B==0 || A%B==0) ans++;
        }
        return ans;
    }
};

これ以外にも中心を決めて左右の重さが等しくなるようなペアを数える方法や，二分探索する方法があるが，累積和の方法が圧倒的に楽． 自分は面倒な方法が実装できなくて 0 pts，hackミスで -25pts を喰らってしまった．