CF 451D - Count Good Substrings

We are given a binary string consisting only of 'a' and 'b'. For every substring, we repeatedly merge consecutive equal characters into a single character.

codeforcescompetitive-programmingmath

CF 451D - Count Good Substrings

Rating: 2000
Tags: math
Solve time: 2m 3s
Verified: yes

Solution

Problem Understanding

We are given a binary string consisting only of 'a' and 'b'. For every substring, we repeatedly merge consecutive equal characters into a single character. For example:

  • "aaaa" becomes "a"
  • "aabbaa" becomes "aba"
  • "abbba" becomes "aba"

A substring is called good if the compressed string obtained after these merges is a palindrome.

The task is not to list the substrings. We only need to count how many good substrings have even length and how many have odd length.

The string length can reach $10^5$. A substring count is already $O(n^2)$, so checking every substring individually is impossible. Even a linear scan per substring would lead to roughly $10^{10}$ operations in the worst case, far beyond what fits in a 2 second time limit. The solution must exploit some hidden structure and run close to linear time.

The most subtle part of the problem is understanding what the compression operation actually changes.

Consider the substring "abba".

After compression:

$$"abba" \rightarrow "aba"$$

which is a palindrome, so the substring is good.

Now consider "abab".

Compression changes nothing:

$$"abab" \rightarrow "abab"$$

and this is not a palindrome.

A common mistake is to think that the entire compressed string must be computed for every substring. The key observation is that for a binary alphabet, the palindrome condition becomes much simpler.

Another easy trap appears with substrings consisting of one repeated character. For example:

aaaa

Every substring compresses to "a", which is a palindrome. Any solution based only on the shape of the compressed string must still count all lengths correctly.

A final edge case is alternating strings such as:

ababab

Compression does nothing, so goodness depends entirely on the original substring. Such cases expose incorrect assumptions about runs of equal characters.

Approaches

A brute force solution would enumerate every substring, compress it, and test whether the compressed result is a palindrome.

There are $O(n^2)$ substrings. A single compression can take $O(n)$ time, giving $O(n^3)$ total complexity. Even with some optimization, checking every substring independently still remains at least $O(n^2)$, which is far too slow for $n=10^5$.

The breakthrough comes from understanding what a good substring looks like in a binary alphabet.

Take any substring. After compressing consecutive equal characters, the result is an alternating sequence because adjacent equal characters have been merged away.

Since only 'a' and 'b' exist, every compressed string must look like one of:

$$a,\quad b,\quad ab,\quad ba,\quad aba,\quad bab,\quad abab,\ldots$$

An alternating binary string is a palindrome exactly when its length is odd.

For example:

$$aba,\ bab,\ ababa$$

are palindromes, while

$$ab,\ ba,\ abab$$

are not.

Now look at a substring $s[l..r]$.

Its compressed form starts with $s[l]$ and ends with $s[r]$.

Since the compressed string alternates, it is a palindrome if and only if its first and last characters are equal.

That means:

A substring is good if and only if its first and last characters are equal.

This completely removes the need for compression.

The problem becomes counting substrings whose endpoints contain the same character, separated into even and odd lengths.

Suppose we process positions from left to right.

For a fixed ending position $i$, every good substring ending at $i$ must start at an earlier position $j$ with:

$$s[j] = s[i]$$

The parity of the substring length is:

$$(i-j+1)$$

Odd length occurs when $i$ and $j$ have the same index parity.

Even length occurs when $i$ and $j$ have opposite parity.

So we only need to know, for each character (a or b), how many previous occurrences appeared at even positions and how many appeared at odd positions.

This gives an $O(n)$ solution.

Approach Time Complexity Space Complexity Verdict
Brute Force $O(n^3)$ $O(n)$ Too slow
Optimal $O(n)$ $O(1)$ Accepted

Algorithm Walkthrough

  1. Maintain four counters.

a_even, a_odd, b_even, and b_odd store how many processed positions contain each character and each index parity. 2. Scan the string from left to right.

Use zero-based indices. 3. For the current position $i$, determine whether $i$ is even or odd. 4. Count odd-length good substrings ending at $i$.

The starting position must contain the same character and have the same parity as $i$.

If the current character is 'a' and $i$ is even, add a_even.

If the current character is 'a' and $i$ is odd, add a_odd.

The same idea applies to 'b'. 5. Count even-length good substrings ending at $i$.

The starting position must contain the same character and have opposite parity.

If the current character is 'a' and $i$ is even, add a_odd.

If the current character is 'a' and $i$ is odd, add a_even.

The same idea applies to 'b'. 6. Count the single-character substring.

Every length-one substring is good and has odd length, so add one to the odd answer. 7. Insert the current position into the corresponding counter.

Future substrings may start here. 8. Continue until the entire string has been processed.

Why it works

A compressed binary string always alternates between 'a' and 'b'. Such a string is a palindrome exactly when its first and last characters are equal. The first and last characters of the compressed string are precisely the first and last characters of the original substring. Hence a substring is good exactly when its endpoints are equal.

After reducing the problem to endpoint equality, parity alone determines whether the substring length is odd or even. Two indices with the same parity produce an odd-length substring, while opposite parities produce an even-length substring. The maintained counters count exactly those earlier positions that can serve as valid starting points, so every good substring is counted once when its right endpoint is processed.

Python Solution

import sys
input = sys.stdin.readline

s = input().strip()

even_good = 0
odd_good = 0

a_even = a_odd = 0
b_even = b_odd = 0

for i, ch in enumerate(s):
    if ch == 'a':
        if i % 2 == 0:
            odd_good += a_even
            even_good += a_odd
            a_even += 1
        else:
            odd_good += a_odd
            even_good += a_even
            a_odd += 1
    else:
        if i % 2 == 0:
            odd_good += b_even
            even_good += b_odd
            b_even += 1
        else:
            odd_good += b_odd
            even_good += b_even
            b_odd += 1

    odd_good += 1

print(even_good, odd_good)

The algorithm processes every position exactly once.

The four counters represent all previously seen occurrences of each character split by index parity. When position i is reached, every good substring ending at i must start at a previous occurrence of the same character.

If the start and end positions have the same parity, the substring length is odd. If their parities differ, the substring length is even. This lets us update both answers in constant time.

The single-character substring at position i is always good, so odd_good is incremented once per character.

The order of operations matters. Contributions from previous positions must be added before inserting the current position into the counters. Otherwise the current position would incorrectly form a length-zero substring with itself.

Python integers automatically handle values up to approximately $n^2$, which is necessary because the answer can reach about $5 \times 10^9$.

Worked Examples

Example 1

Input:

bb
i char parity even_good odd_good b_even b_odd
0 b even 0 1 1 0
1 b odd 1 2 1 1

Final answer:

1 2

The two single-character substrings contribute to the odd count. The substring "bb" has equal endpoints and even length, so it contributes to the even count.

Example 2

Input:

baab
i char parity even_good odd_good a_even a_odd b_even b_odd
0 b even 0 1 0 0 1 0
1 a odd 0 2 0 1 1 0
2 a even 1 3 1 1 1 0
3 b odd 2 4 1 1 1 1

Final answer:

2 4

The even substrings are "aa" and "baab". The four single-character substrings form the odd count.

Complexity Analysis

Measure Complexity Explanation
Time $O(n)$ One pass through the string
Space $O(1)$ Only a few counters are stored

The string length reaches $10^5$, so a linear scan performs only $10^5$ iterations. Constant extra memory easily fits within the memory limit. Both constraints are comfortably satisfied.

Test Cases

# helper: run solution on input string, return output string
import sys, io

def solve():
    s = input().strip()

    even_good = 0
    odd_good = 0

    a_even = a_odd = 0
    b_even = b_odd = 0

    for i, ch in enumerate(s):
        if ch == 'a':
            if i % 2 == 0:
                odd_good += a_even
                even_good += a_odd
                a_even += 1
            else:
                odd_good += a_odd
                even_good += a_even
                a_odd += 1
        else:
            if i % 2 == 0:
                odd_good += b_even
                even_good += b_odd
                b_even += 1
            else:
                odd_good += b_odd
                even_good += b_even
                b_odd += 1

        odd_good += 1

    print(even_good, odd_good)

def run(inp: str) -> str:
    sys.stdin = io.StringIO(inp)
    global input
    input = sys.stdin.readline

    out = io.StringIO()
    old_stdout = sys.stdout
    sys.stdout = out

    solve()

    sys.stdout = old_stdout
    return out.getvalue().strip()

# provided sample
assert run("bb\n") == "1 2", "sample 1"

# minimum size
assert run("a\n") == "0 1", "single character"

# sample 2 from statement
assert run("baab\n") == "2 4", "statement example"

# alternating string
assert run("abab\n") == "0 6", "alternating pattern"

# all equal
assert run("aaaa\n") == "4 6", "all substrings good"

# off-by-one parity check
assert run("aba\n") == "1 4", "same endpoints at distance two"
Test input Expected output What it validates
a 0 1 Minimum size
baab 2 4 Mixed runs and compression
abab 0 6 Alternating pattern
aaaa 4 6 Every substring is good
aba 1 4 Parity handling and endpoint counting

Edge Cases

Consider the input:

aaaa

Every substring compresses to "a", so every substring is good. The algorithm counts substrings purely through endpoint equality. Since every endpoint pair contains 'a', every substring is counted. The result is:

4 6

which matches the ten total substrings of a length four string.

Now consider:

abab

Compression changes nothing. Good substrings are exactly those whose first and last characters match. The only such substrings are the four single letters and the two length-three substrings. The algorithm reaches:

0 6

because no even-length substring has matching endpoints.

Finally consider:

abba

The whole string compresses to "aba", which is a palindrome. A naive implementation that reasons only about the original string might miss this. The endpoint characterization handles it immediately: the first and last characters are both 'a', so the substring is good. The algorithm counts it through the matching 'a' occurrences at positions $0$ and $3$, adding one even-length good substring.