Medium LC Problem in Python

This is my first medium-difficulty LC problem in Python (learning the language). It still does not require any deep understand of the language (only simple data structures), just the basics loops and conditions. Here is the problem: https://leetcode.com/problems/number-of-ways-to-split-a-string/

1573. Number of Ways to Split a String

Medium

Given a binary string s (a string consisting only of '0's and '1's), we can split s into 3 non-empty strings s1, s2, s3 (s1+ s2+ s3 = s).

Return the number of ways s can be split such that the number of characters '1' is the same in s1, s2, and s3.

Since the answer may be too large, return it modulo 10^9 + 7.

 

Example 1:

Input: s = "10101"
Output: 4
Explanation: There are four ways to split s in 3 parts where each part contain the same number of letters '1'.
"1|010|1"
"1|01|01"
"10|10|1"
"10|1|01"

Example 2:

Input: s = "1001"
Output: 0

Example 3:

Input: s = "0000"
Output: 3
Explanation: There are three ways to split s in 3 parts.
"0|0|00"
"0|00|0"
"00|0|0"

Example 4:

Input: s = "100100010100110"
Output: 12

 

Constraints:

• s[i] == '0' or s[i] == '1'
• 3 <= s.length <= 10^5

The logic is actually pretty straightforward after some initial investigation:
 - If the string length is less than 3, return 0
 - If it does not contain 1s, just do a combinatorial
 - Else, find the zeros in between the blocks, add one, and multiply the result, always "mod"ing it

Code is below, cheers, ACC.

    def numWays(self, s):
        #count the number of 1s and return it
        numberOnes = 0
        for i in range(len(s)):
            if s[i] == '1':
                numberOnes += 1
        
        #if zero, simple combinatorial
        if numberOnes == 0:
            if len(s) < 3:
                return 0
            return int((((len(s) - 1)*(len(s) - 2))/2) % 1000000007)

        #not multiple of 3, bail
        if numberOnes % 3 != 0:
            return 0

        #solution is the multiplication of number of zeros (plus 1) in between the blocks
        index1 = numberOnes/3
        countBlock1 = 0
        insideBlock1 = False
        index2 = 2*index1
        countBlock2 = 0
        insideBlock2 = False

        countOnes = 0
        for i in range(len(s)):
            if s[i] == '1':
                countOnes += 1
            if not insideBlock1 and countOnes == index1:
                insideBlock1 = True
                continue
            if countOnes == index1 + 1:
                insideBlock1 = False
            if not insideBlock2 and countOnes == index2:
                insideBlock2 = True
                continue
            if countOnes == index2 + 1:
                insideBlock2 = False

            if insideBlock1:
                countBlock1 += 1
            if insideBlock2:
                countBlock2 += 1

        return int(((countBlock1 + 1) * (countBlock2 + 1)) % 1000000007)