BinarySearch 0911 Three-Way String Split with Equal Ones

Problem statement

https://binarysearch.com/problems/Three-Way-String-Split-with-Equal-Ones/

Solution

Equal to Leetcode 1573 Number of Ways to Split a String.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def solve(self, s):
        M = 10**9 + 7
        n, k = len(s), s.count("1")
        if k % 3 != 0: return 0
        if k == 0: return (n-1)*(n-2)//2 % M
        P = [i for i in range(n) if s[i] == "1"]
        return (P[k//3] - P[k//3 - 1]) * (P[2*k//3] - P[2*k//3 - 1]) % M