TCO15 Round 2D easy:BalancedSubstrings

TCO15 Round 2D に参加しました。easy 解くのが遅すぎてひどかったですが, チャレンジを 1 回成功させたおかげでレート微増してくれました(challenge のありがたみをはじめて感じた)。もちろん Round 3 には進めませんでした。

問題

TopCoder Statistics - Problem Statement

解法

自分はで解きました。

左と右を決めて真ん中の座標を二分探索する。
そのために memo[i] = (0 <= x <= i にあるs[x] == '1' となる x の数), sum[i] = (0 <= x <= i にあるs[x] == '1' となる x の座標値の和)を前計算しておくと, 中央から左端までのトルク, 中央から右端までのトルクが O(1) で計算できる。

const int MAXN = 2505;
int memo[MAXN], sum[MAXN];

int query(int l, int m, int r) {
    int lnum = memo[m-1] - memo[l-1];
    int rnum = memo[r] - memo[m];
    int left = m * lnum - (sum[m-1]-sum[l-1]);
    int right = sum[r]-sum[m] - rnum*m;
    return right-left;
}

class BalancedSubstrings {
public:
    int countSubstrings(string s) {
        cout << "TEST" << endl;
        int n = s.size();
        memo[0] = 0;
        sum[0] = 0;
        for (int i = 1; i <= n; i++) {
            memo[i] = memo[i-1];
            sum[i] = sum[i-1];
            if (s[i-1] == '1') {
                memo[i]++;
                sum[i] += i;
            }
        }
        int ans = 0;
        for (int l = 1; l <= n; l++) for (int r = l; r <= n; r++) {
            int low = l, high = r+1;
            while (high-low > 1) {
                int med = (high+low)/2;
                if (query(l, med, r) >= 0) low = med;
                else high = med;
            }
            if (query(l, low, r) == 0) ans++;
        }
        return ans;
    }
};