题解 | 小红的01子序列计数（hard）

#include <iostream>
#include <string>
using namespace std;

#define rep(i, a, b) for (int i = (a), _##i = (b); i <= _##i; ++i)

using ll = long long;

const int N = 1e5 + 5;
const int mod = 1e9 + 7;

// 全局变量
ll ans = 0;    // 最终答案
ll sum = 0;    // 当前窗口内0的数量，用于辅助更新
ll num0 = 0;   // 当前窗口内0的数量
ll sum0 = 0;   // 当前窗口内所有0的下标之和
ll num1 = 0;   // 当前窗口内1的数量
ll sum1 = 0;   // 当前窗口内"01"子序列数量

void solve() {
    int n;
    cin >> n;
    string s;
    cin >> s;

    // 为了模拟环形，把s复制一份自己接到自己后面，并在前面加空格使下标从1开始
    s = " " + s + s;

    // 初始化滑动窗口，处理前n个字符
    rep(i, 1, n) {
        if (s[i] == '0') {
            num0++;
            sum0 += i;
        } else { // s[i] == '1'
            num1++;
            sum1 += sum0;  // 新加入的1，与之前所有0形成"01"子序列
            sum += num0;   // 统计出现过的0的数量
        }
    }

    // 滑动窗口，处理后续n个字符
    rep(i, n + 1, 2 * n) {
        // 移除窗口左端点元素 i - n 的影响
        sum1 -= sum;      // 去掉以移除元素为起点产生的贡献
        sum0 -= num0;     // 更新所有0下标和
        if (s[i - n] == '0') {
            num0--;       // 0数量减少
            sum -= num1;  // 0被移除，减少对后续1的贡献
        } else {
            num1--;       // 1数量减少
            // 移除1不会影响sum
        }

        // 加入新的元素 s[i]
        if (s[i] == '0') {
            num0++;
            sum0 += n;    // 加入0，位置视为n（窗口固定大小）
        } else { // s[i] == '1'
            num1++;
            sum1 += sum0; // 新加入的1，与已有0形成新子序列
            sum += num0;  // 0的数量影响sum
        }

        // 累加当前窗口的贡献
        ans = (ans + sum1) % mod;
    }

    cout << ans << '\n';
}

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

    int t = 1;
    // cin >> t; // 如果有多组测试数据可以打开

    while (t--) {
        solve();
    }

    return 0;
}