OZ1 - Editorial

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

None

PROBLEM:

You’re given a binary string S.
A prefix of S is called good if it contains at least as many ones as zeros.

Count the number of good prefixes of S.

EXPLANATION:

We simply need to check, for each prefix of S, if it contains at least as many zeros as ones.

This can be done easily by maintaining a running count of zeros and ones.
Specifically,

  • Let \text{Z} and O denote the number of zeros and ones in the current prefix.
    Initially, these values are both 0.
  • For each i = 1, 2, 3, \ldots, N in order,
    • If S_i = \texttt{0}, add 1 to Z.
    • Otherwise, add 1 to O.
    • Then, if O \ge Z, add 1 to the answer.
  • In the end, print the answer.

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

Editorialist's code (PyPy3)
for _ in range(int(input())):
    n = int(input())
    s = input()
    
    zeros, ones = 0, 0
    ans = 0
    for i in range(n):
        if s[i] == '0': zeros += 1
        else: ones += 1
        if ones >= zeros: ans += 1
    print(ans)

This question have wrong statement on output

@iceknight1093 Why does it say this for test case 2 in the problem statement,

Test Case 1, 2 : There are no good prefixes.

4
3
000
3
001
5
01110
5
11111

If you do a swap in 001 and obtain 010, you will have a binary prefix 01 that does satsify the condition: “A prefix S[1,i] is called good if it contains at least as many ones as zeroes.”?

Swapping is not a thing in OZ1 (Ones and Zeroes I).
The testcases are correct.