【Leetcode】1573. Number of Ways to Split a String

题目地址：

https://leetcode.com/problems/number-of-ways-to-split-a-string/description/

给定一个非空长$n$的$01$字符串$s$，欲将其分为$3$个非空子串，使得这$3$个子串含$1$的个数相等。问方案数。

统计一下$s$中一共多少个$1$，假设有$c$个$1$，如果$3\nmid c$那显然无解，返回$0$；如果$c=0$，则问题相当于问在$n-1$个空里选两个空有多少个方法，方案数为$\frac{(n-1)(n-2)}{2}$；如果$c>0$，那么我们需要找到两端的两个字符串的边界的位置，设$s[0:c_1]$是最短的含$c/3$个$1$的前缀，设其后紧跟着$k_1$个$0$，再设$s[c_2:]$是最短的含$c/3$个$1$的后缀，设其前紧跟着$k_2$个$0$，则两个隔板需要放在这两段$0$里，每段可以放的位置是$k_1+1,k_2+1$，所以方案数就是$(k_1+1)(k_2+1)$。代码如下：

class Solution {
 public:
  int numWays(string s) {
    const int MOD = 1e9 + 7;
    int n = s.size(), cnt = 0;
    for (char &ch : s) cnt += ch - '0';
    if (cnt % 3) return 0;

    if (!cnt) return ((long)(n - 1) * (n - 2) / 2) % MOD;

    cnt /= 3;
    int x[2];
    for (int i = 0, c = 0; i < n; i++) {
      c += s[i] - '0';
      if (c == cnt || c == cnt * 2) {
        int j = i + 1;
        while (j < n && s[j] == '0') j++;
        x[c / cnt - 1] = j - i - 1;
        i = j - 1;
      }
    }

    return ((long)(x[0] + 1) * (x[1] + 1) % MOD);
  }
};

时间复杂度$O(n)$，空间$O(1)$。