3003. Maximize the Number of Partitions After Operations

HardBit ManipulationStringDynamic ProgrammingBitmask
Leetcode Link

Problem Description

The problem presents a string s that is indexed starting from 0, along with an integer k. The goal is to repeatedly partition the string into prefixes that contain at most k distinct characters. After each partition, the selected prefix is deleted from the string and the count of partitions is increased by one. Before beginning the partitioning operations, there is an option to change one character in the string to any other lowercase English letter. The task is to determine the maximum number of partitions that can be obtained by optimally choosing to change at most one character.

Intuition

To achieve the most number of partitions, we can use a recursive approach with a depth-first search (DFS) strategy. The intuition behind the solution is based on the fact that as we traverse the string, we choose the longest prefix with at most k distinct characters. If we encounter more than k distinct characters, we need to either start a new partition or consider changing one character if that option hasn't been used yet.

Since trying all possible character changes would be inefficient, we use memoization (caching) to remember the results of computations for particular states to avoid redundant work. This is evident in the use of the @cache decorator in the provided solution code.

The recursive function dfs keeps track of the current index in the string, the bitmask cur representing the set of distinct characters seen so far, and the t flag indicating whether we can still change a character. The bitmask is particularly useful as it allows us to quickly count the number of distinct characters in the current prefix and to efficiently apply a character change by updating the bitmask. The DFS is designed to explore two scenarios at each position: continue increasing the current prefix, or start a new partition. This helps us to find the optimal partitions that lead to the maximum count.

The recursion stops when we have reached the end of the string, and at each step, we make choices based on whether it's better to change a character or to split the partition at the current point if the number of distinct characters goes over k.

Utilizing this approach, leveraging bit manipulation and memoization, we can explore the space efficiently and arrive at the solution that provides the maximum number of partitions.

Learn more about Dynamic Programming and Bitmask patterns.

Solution Approach

The algorithm begins by defining a recursive function dfs that employs Depth-First Search (DFS). This dfs function takes three arguments: the current index i, the bitmask cur representing the distinct characters currently in the prefix, and the t flag to indicate whether the character change option is still available. The function returns an integer indicating the number of partitions formed from the current index onward.

The core idea lies in the use of memoization, where the @cache decorator stores the results of recursive calls, thereby avoiding repeated calculations for the same arguments. This optimization is significant as it reduces the complexity by ensuring each state is computed only once.

Inside the DFS, we perform the following steps:

  1. Check if we have processed the entire string. If i >= n, where n is the length of the string, we return 1 as no further partition can be made.

  2. Calculate a bitmask v for the current character. This is done using 1 << (ord(s[i]) - ord("a")) to find the bit corresponding to the current character in the alphabet.

  3. Determine the new bitmask nxt after adding the current character to the prefix by performing a bitwise OR operation (|) with the bitmask cur.

  4. If the new bitmask nxt has more than k distinct characters (nxt.bit_count() > k), we've reached a prefix with too many distinct characters, meaning we must partition here. In this case, we recurse with index i + 1, start a new bitmask with just the new character (v), and add 1 to the partition count.

  5. If nxt still has k or fewer distinct characters, we continue with the current partition and recurse with i + 1 and nxt without changing the partition count.

  6. If we still have the option to change a character (t is true), we iterate over all possible 26 lowercase English letters. For each, we calculate the bitmask nxt as if we changed the current character to each option.

  7. Depending on whether this hypothetical change results in a prefix with too many distinct characters, we recurse either with the same logic as in step 4 (and thus a new partition) or continue the current prefix with the modified bitmask (step 5), always selecting the option that yields the maximum number of partitions.

The n (string length) is calculated once before invoking the dfs function, and the function is first called with initial arguments dfs(0, 0, 1) (i.e., starting at index 0, with an empty bitmask, and 1 indicating we have the option to change one character).

The complexity involves the length of the string n and the maximum possible bitmask value, which corresponds to the number of distinct lowercase English letters 2^26. The memoization effectively reduces time complexity by avoiding repeating calculations for the same state.

The use of bit operation .bit_count(), a method introduced in Python 3.10, quickly counts the number of bits set to 1 in the bitmask, which represents the count of distinct characters currently considered in the prefix.

The result from the first call to dfs gives the maximum number of resulting partitions possible after performing the optimal character change and partitioning operations.

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Which data structure is used to implement recursion?

Example Walkthrough

Consider a small example to illustrate the solution approach. Let's take the string s = "aabbaa" and k = 2. Our target is to make the maximum number of partitions such that each contains at most 2 distinct characters. We also have the option to change one character to try to increase the total number of partitions.

Initial call to DFS: dfs(0, 0, 1)

Here's how the function executes:

  • We start at index 0 with an empty bitmask (which means no characters have been found yet), and the t flag is true (meaning we have the option to change a character).

  • Moving to the first character 'a', we update the bitmask cur to include it. The bitmask v for 'a' will be 1 << (ord('a') - ord('a')) = 1. Then nxt becomes 0 | 1 = 1. As the number of distinct characters (nxt.bit_count()) is still within the limit of k, we can opt to continue without partitioning.

  • We do the same for the second 'a' at index 1 and the first 'b' at index 2. The bitmask v for 'b' is 1 << (ord('b') - ord('a')) = 10 in binary, and thus nxt is now 1 | 10 = 11 in binary. We have encountered two distinct characters, 'a' and 'b', but we are still within the limit.

  • At index 3, we also append the 'b' to our current prefix. nxt remains 11 in binary because 'b' is already in our bitmask.

  • When we reach index 4 and see another 'a', appending it to the current prefix doesn't change the distinct characters. We could continue without partitioning and proceed to index 5.

  • At index 5, appending 'a' again doesn't introduce any new characters. Since we reached the end of the string, we can consider this a valid partition, so we add 1 to partition count.

However, we may also consider using our character change option at certain points. For example:

  • At index 4, we could have changed 'a' to 'c' (if we hadn't used our change yet), which would cause our nxt.bit_count() to exceed k, hence we would have had to create a new partition there.

For this example, we'll find that we don't actually need to use the character change option to maximize the partitions; we can achieve two partitions with the original characters: "aabba" and "a".

But, if k was 1, we could consider changing the second 'b' to an 'a' to create a partition of "aabaa" and then have the final 'a' as a partition on its own, which would give us two partitions in total, maximizing the number of partitions given the constraints.

This recursive approach, combined with memoization, would try these operations and eventually yield the maximum number of partitions possible. In the case of our s = "aabbaa" and k = 2, this maximum number would be 2.

Solution Implementation

1from functools import cache
2
3class Solution:
4    def maxPartitionsAfterOperations(self, s: str, k: int) -> int:
5        # Helper function to recursively determine the maximum number of partitions
6        # i: current position in string
7        # current_mask: the bit mask representing the current partition's character set
8        # t: the tracker if second type of operations are still available
9        @cache
10        def dfs(i: int, current_mask: int, t: int) -> int:
11            # Base case: if we've reached the end of the string return 1 partition
12            if i >= n:
13                return 1
14          
15            # Calculate the bit value for the current character
16            char_bit = 1 << (ord(s[i]) - ord("a"))
17            # Combine the current character bit with the mask for the next state
18            next_mask = current_mask | char_bit
19          
20            # Check if combining this character exceeds the k unique characters limit
21            if next_mask.bit_count() > k:
22                # If it does, partition here and reset mask for new partition
23                ans = dfs(i + 1, char_bit, t) + 1
24            else:
25                # Otherwise, continue with the current partition
26                ans = dfs(i + 1, next_mask, t)
27          
28            # If the second type of operations are still available
29            if t:
30                # Go through all possible characters and generate new masks
31                for j in range(26):
32                    # Create a mask with only this character
33                    next_mask = current_mask | (1 << j)
34                    # Again, check if this mask exceeds the unique characters limit
35                    if next_mask.bit_count() > k:
36                        # Recurse with a new partition for this character
37                        ans = max(ans, dfs(i + 1, 1 << j, 0) + 1)
38                    else:
39                        # Recurse with an updated partition including this character
40                        ans = max(ans, dfs(i + 1, next_mask, 0))
41          
42            # Return the maximum partitions found
43            return ans
44
45        # Length of the string
46        n = len(s)
47      
48        # Start the recursive DFS at position 0, with no characters in the partition, and operations available
49        return dfs(0, 0, 1)
50
1import java.util.HashMap;
2import java.util.List;
3import java.util.Map;
4
5class Solution {
6    private Map<List<Integer>, Integer> memoization = new HashMap<>();
7    private String inputString;
8    private int maxUniqueLetters;
9
10    // Entry method to start the process with given string 's' and 'k' maximum unique characters
11    public int maxPartitionsAfterOperations(String s, int k) {
12        this.inputString = s;
13        this.maxUniqueLetters = k;
14        return depthFirstSearch(0, 0, 1);
15    }
16
17    // Recursive DFS method to find the maximum number of partitions
18    private int depthFirstSearch(int index, int current, int toggle) {
19        // Base case: if we have gone past the last character
20        if (index >= inputString.length()) {
21            return 1;
22        }
23        // Cache key made of current state
24        var key = List.of(index, current, toggle);
25        // If the result for the current state is already computed, return it
26        if (memoization.containsKey(key)) {
27            return memoization.get(key);
28        }
29        // Compute the bit representation of the current character
30        int charBitValue = 1 << (inputString.charAt(index) - 'a');
31        // OR operation to include the current character in the set
32        int next = current | charBitValue;
33        // Compute the answer recursively based on the number of unique characters
34        int answer = Integer.bitCount(next) > maxUniqueLetters ? 
35                        depthFirstSearch(index + 1, charBitValue, toggle) + 1 : 
36                        depthFirstSearch(index + 1, next, toggle);
37        // If toggling is allowed, try to optimize answer by considering more partitions
38        if (toggle > 0) {
39            for (int j = 0; j < 26; ++j) {
40                next = current | (1 << j); // Toggle each character
41                if (Integer.bitCount(next) > maxUniqueLetters) {
42                    // If new set exceeds max unique chars limit, partition here
43                    answer = Math.max(answer, depthFirstSearch(index + 1, 1 << j, 0) + 1);
44                } else {
45                    // Otherwise, continue with the same set
46                    answer = Math.max(answer, depthFirstSearch(index + 1, next, 0));
47                }
48            }
49        }
50        // Cache the result for current state
51        memoization.put(key, answer);
52        return answer;
53    }
54}
55
1#include <unordered_map>
2#include <functional>
3#include <string>
4
5class Solution {
6public:
7    // Defines a method that calculates the maximum number of partitions 
8    // after specific operations can be performed on the string.
9    int maxPartitionsAfterOperations(std::string s, int k) {
10        int n = s.size(); // Length of the input string
11        std::unordered_map<long long, int> memo; // To memoize intermediate results
12
13        // Defines a lambda function for depth-first search.
14        std::function<int(int, int, int)> dfs = [&](int index, int bitMask, int canChange) {
15            // Base case: when we have considered all characters in the string
16            if (index >= n) {
17                return 1;
18            }
19            // Create a unique key for memorization from the current state
20            long long key = (long long)index << 32 | bitMask << 1 | canChange;
21            // Check if the result for this state was already calculated
22            if (memo.count(key)) {
23                return memo[key];
24            }
25            int characterBit = 1 << (s[index] - 'a'); // Bit representation of the current character
26            int nextBitMask = bitMask | characterBit; // Next state bit mask with the current character included
27            // Calculate the answer for the current state through recursion
28            int ans = __builtin_popcount(nextBitMask) > k ? dfs(index + 1, characterBit, canChange) + 1 : dfs(index + 1, nextBitMask, canChange);
29          
30            // If we are allowed to change characters
31            if (canChange) {
32                // Try changing the current character to all other possible characters (a-z)
33                for (int j = 0; j < 26; ++j) {
34                    nextBitMask = bitMask | (1 << j);
35                    // Recursively call dfs with the new character, considering the k limit
36                    if (__builtin_popcount(nextBitMask) > k) {
37                        ans = std::max(ans, dfs(index + 1, 1 << j, 0) + 1);
38                    } else {
39                        ans = std::max(ans, dfs(index + 1, nextBitMask, 0));
40                    }
41                }
42            }
43            // Memoize and return the answer for the current state
44            return memo[key] = ans;
45        };
46
47        // Start the DFS from the beginning of the string
48        // 0 as the initial bitmask and 1 to allow for changing characters initially
49        return dfs(0, 0, 1);
50    }
51};
52
1/**
2 * Calculates the maximum number of partitions possible after performing
3 * transformation operations under specific rules.
4 * @param s The input string composed of lowercase English letters.
5 * @param k The maximum number of unique characters that can appear in one partition.
6 * @returns The maximum number of partitions after transformation.
7 */
8function maxPartitionsAfterOperations(s: string, k: number): number {
9    const length = s.length;
10    const memoization: Map<bigint, number> = new Map();
11  
12    /**
13     * Depth-first search helper function to recursively determine the maximum partitions.
14     * @param index The current index in the string.
15     * @param currentMask The bitmask representing the current partition's unique characters.
16     * @param useTransformation Indicates whether a transformation is used at this step.
17     * @returns The maximum number of partitions from this point.
18     */
19    const dfs = (index: number, currentMask: number, useTransformation: number): number => {
20        if (index >= length) {
21            return 1;
22        }
23        const key = (BigInt(index) << 27n) | (BigInt(currentMask) << 1n) | BigInt(useTransformation);
24        if (memoization.has(key)) {
25            return memoization.get(key)!;
26        }
27        const currentCharMask = 1 << (s.charCodeAt(index) - 'a'.charCodeAt(0));
28        let nextMask = currentMask | currentCharMask;
29        let answer = 0;
30      
31        // If adding the current character exceeds the number of unique characters allowed,
32        // it starts a new partition.
33        if (bitCount(nextMask) > k) {
34            answer = dfs(index + 1, currentCharMask, useTransformation) + 1;
35        } else {
36            answer = dfs(index + 1, nextMask, useTransformation);
37        }
38
39        // If transformations can still be used, try transforming the current character
40        // to each possible character and see which gives the most partitions.
41        if (useTransformation) {
42            for (let charIndex = 0; charIndex < 26; ++charIndex) {
43                nextMask = currentMask | (1 << charIndex);
44                if (bitCount(nextMask) > k) {
45                    answer = Math.max(answer, dfs(index + 1, 1 << charIndex, 0) + 1);
46                } else {
47                    answer = Math.max(answer, dfs(index + 1, nextMask, 0));
48                }
49            }
50        }
51        memoization.set(key, answer);
52        return answer;
53    };
54
55    // Start the depth-first search with an empty partition, using transformation.
56    return dfs(0, 0, 1);
57}
58
59/**
60 * Counts the number of bits set to 1 in a 32 bit integer.
61 * @param number The number to count bits in.
62 * @returns The number of bits set to 1.
63 */
64function bitCount(number: number): number {
65    number = number - ((number >>> 1) & 0x55555555);
66    number = (number & 0x33333333) + ((number >>> 2) & 0x33333333);
67    number = (number + (number >>> 4)) & 0x0f0f0f0f;
68    number = number + (number >>> 8);
69    number = number + (number >>> 16);
70    return number & 0x3f;
71}
72

Time and Space Complexity

Time Complexity

The time complexity of the algorithm is predominantly determined by the dfs function, which is a recursive function with memoization. The dfs function takes three arguments – i, cur, and t. The first, i, ranges from 0 to n, where n is the length of the input string s. The second, cur, which represents the current state, can have at most 2^k different values because when the bit count exceeds k, the partition happens. The third, t, can be either 0 or 1. Therefore, the state space for memoization is O(2^k * n * 2).

Within each recursive call, the function checks the bit count of the current state and potentially performs two recursive calls. One call is to continue adding letters to the current partition, and the other is to start a new partition if the bit count limit is exceeded. Additionally, when t is 1 (which can happen once at the beginning for each index i), there is a loop running for the 26 alphabet letters that could potentially make 26 more recursive calls.

Assuming bit_count() is O(1), under the hood each recursive call contributes a constant amount of work plus the work from further recursive calls.

The worst-case time complexity (without considering memoization) could be as high as O(26 * n * 2^k), assuming each recursive call will iterate over the 26 letters when t is 1. But with memoization, we do not recompute states that we have already computed, which substantially reduces the actual number of recursive calls performed.

Overall, the time complexity is O(26 * n * 2^k) before memoization and O(2^k * n * 2) with memoization, as each state would be computed at most once.

Space Complexity

The space complexity is primarily due to the memoization cache and the recursion stack.

The cache will store every unique state that is computed. Since the state is based on i, cur, and t, the maximum number of entries in the cache will be O(2^k * n * 2), as discussed above for the time complexity. Also, the recursion stack will go as deep as the string length n, which contributes O(n) space in the worst case.

Therefore, the overall space complexity is O(2^k * n * 2 + n). Since 2^k * n * 2 is likely the dominant term, we can simplify this to O(2^k * n).

Learn more about how to find time and space complexity quickly using problem constraints.


Fast Track Your Learning with Our Quick Skills Quiz:

What's the output of running the following function using the following tree as input?

1def serialize(root):
2    res = []
3    def dfs(root):
4        if not root:
5            res.append('x')
6            return
7        res.append(root.val)
8        dfs(root.left)
9        dfs(root.right)
10    dfs(root)
11    return ' '.join(res)
12
1import java.util.StringJoiner;
2
3public static String serialize(Node root) {
4    StringJoiner res = new StringJoiner(" ");
5    serializeDFS(root, res);
6    return res.toString();
7}
8
9private static void serializeDFS(Node root, StringJoiner result) {
10    if (root == null) {
11        result.add("x");
12        return;
13    }
14    result.add(Integer.toString(root.val));
15    serializeDFS(root.left, result);
16    serializeDFS(root.right, result);
17}
18
1function serialize(root) {
2    let res = [];
3    serialize_dfs(root, res);
4    return res.join(" ");
5}
6
7function serialize_dfs(root, res) {
8    if (!root) {
9        res.push("x");
10        return;
11    }
12    res.push(root.val);
13    serialize_dfs(root.left, res);
14    serialize_dfs(root.right, res);
15}
16

Recommended Readings


Got a question? Ask the Monster Assistant anything you don't understand.

Still not clear? Ask in the Forum,  Discord or Submit the part you don't understand to our editors.


🪄