Topics › Strings › Medium String Problems

Count Number of Substrings

Easy TUF ID 2394 A2Z

Problem

You are given a string s of lowercase English letters and an integer k. Count the substrings of s that contain exactly k distinct characters. Substrings are contiguous, and two substrings that have the same letters but start or end at different positions count separately.

Constraints: 1 ≤ |s| ≤ 10^5, 1 ≤ k ≤ 26. The answer can be as large as n(n + 1) / 2 ≈ 5 × 10^9, which does not fit in an int, so the methods return long.

Example 1

s = "aba", k = 2
answer = 3              // "ab", "ba", "aba"

Example 2

s = "abaaca", k = 1
answer = 7              // "a", "b", "a", "a", "aa", "c", "a"

Example 3

s = "aaaa", k = 1
answer = 10             // every one of the n(n+1)/2 = 10 substrings qualifies

Approach 1: Brute force, check every substring

Intuition. Enumerate every pair (i, j) with i ≤ j, count the distinct letters of s[i..j] with a 26-slot boolean table, and increment the answer when that count equals k. This is the definition of the problem written as code; nothing is shared between substrings, so it does a lot of repeated work.

Algorithm.

  1. For every start i and every end j ≥ i:
  2. Clear a boolean[26], walk p from i to j, and count the letters seen for the first time.
  3. If the count equals k, add 1 to the answer.
static long bruteForce(String s, int k) {
    int n = s.length();
    long count = 0;
    for (int i = 0; i < n; i++) {
        for (int j = i; j < n; j++) {
            boolean[] seen = new boolean[26];
            int distinct = 0;
            for (int p = i; p <= j; p++) {
                int c = s.charAt(p) - 'a';
                if (!seen[c]) { seen[c] = true; distinct++; }
            }
            if (distinct == k) count++;
        }
    }
    return count;
}

Complexity. Time O(n^3): O(n^2) substrings, each scanned in O(n). Space O(1): the 26-entry table.

Approach 2: Better, fix the start and grow the end

Intuition. All substrings that begin at the same index i form a chain: s[i..i], s[i..i+1], s[i..i+2], and so on. Each one is the previous one plus a single letter. So instead of recounting, keep a frequency table for the current s[i..j] and update it when j advances: a letter whose count goes from 0 to 1 raises the distinct count by one. This drops a factor of n.

There is one more thing to notice: adding letters can never reduce the number of distinct characters. Once the chain reaches k + 1 distinct letters, every longer substring starting at i also has more than k, so we can break and move to the next start.

Algorithm.

  1. For every start i: clear freq[26] and set distinct = 0.
  2. For j from i to n - 1: increment freq[s[j]]; if it was 0 before, distinct++.
  3. If distinct == k, count this substring. If distinct > k, break out of the inner loop.
static long better(String s, int k) {
    int n = s.length();
    long count = 0;
    for (int i = 0; i < n; i++) {
        int[] freq = new int[26];
        int distinct = 0;
        for (int j = i; j < n; j++) {
            int c = s.charAt(j) - 'a';
            if (freq[c]++ == 0) distinct++;
            if (distinct == k) count++;
            else if (distinct > k) break;              // extending further can only add distinct letters
        }
    }
    return count;
}

Complexity. Time O(n^2) in the worst case, for instance when k is at least the number of distinct letters in s so the break never fires. Space O(1).

For n = 10^5 this is about 5 × 10^9 steps in the worst case, so a linear method is needed.

Approach 3: Optimal, exactly k = at most k minus at most k - 1

Intuition. Counting substrings with exactly k distinct letters is awkward for a sliding window, because the set of valid start positions for a fixed end is a range in the middle of the string, bounded on both sides. Counting substrings with at most k distinct letters is easy, because the valid starts for a fixed end form a suffix of the start positions: if s[left..right] has at most k distinct letters, then so does every s[l..right] with l ≥ left, since it contains a subset of those letters.

Then use the identity

exactly(k) = atMost(k) - atMost(k - 1)

which holds because every substring with at most k distinct letters has either exactly k or at most k - 1, and these two groups do not overlap.

To compute atMost(k) in linear time, slide a window [left, right]. Move right one step at a time, add s[right] to the frequency table, and while the window has more than k distinct letters remove s[left] and advance left. After the shrink, the window s[left..right] is the longest valid window ending at right, and the substrings ending at right that are valid are exactly those starting at left, left + 1, ..., right: there are right - left + 1 of them. Add that to the count.

Why is this linear? right moves n times and left only ever moves forward, so it also moves at most n times in total.

Algorithm.

  1. If k ≤ 0 return 0 (no non-empty substring has zero distinct letters).
  2. Return atMost(s, k) - atMost(s, k - 1).
  3. atMost(s, k): for each right, add s[right] to the table; while distinct > k, remove s[left] and advance left; add right - left + 1 to the count.
static long optimal(String s, int k) {
    if (k <= 0) return 0;
    return atMost(s, k) - atMost(s, k - 1);
}
private static long atMost(String s, int k) {
    if (k <= 0) return 0;
    int[] freq = new int[26];
    int distinct = 0;
    long count = 0;
    for (int left = 0, right = 0; right < s.length(); right++) {
        if (freq[s.charAt(right) - 'a']++ == 0) distinct++;
        while (distinct > k) {                         // shrink until the window is valid again
            if (--freq[s.charAt(left) - 'a'] == 0) distinct--;
            left++;
        }
        count += right - left + 1;                     // every substring ending at right that starts in [left, right]
    }
    return count;
}

Dry run of atMost on s = "aba" for k = 2 and k = 1:

right letter k = 2: window after shrink added k = 1: window after shrink added
0 a [0, 0] “a” 1 [0, 0] “a” 1
1 b [0, 1] “ab” 2 [1, 1] “b” 1
2 a [0, 2] “aba” 3 [2, 2] “a” 1

atMost(2) = 6, atMost(1) = 3, so exactly(2) = 3, matching Example 1.

Complexity. Time O(n): two sliding-window passes, each with at most 2n pointer moves. Space O(1): two 26-entry tables.

Edge cases and pitfalls