Skip to content

3003. Maximize the Number of Partitions After Operations 👍

  • Time: $O(26n) = O(n)$
  • Space: $O(26n) = O(n)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class Solution {
 public:
  int maxPartitionsAfterOperations(string s, int k) {
    unordered_map<long, int> mem;
    return maxPartitionsAfterOperations(s, 0, true, 0, k, mem) + 1;
  }

 private:
  // Returns the maximum number of partitions of s[i..n), where `canChange` is
  // true if we can still change a letter, and `mask` is the bitmask of the
  // letters we've seen.
  int maxPartitionsAfterOperations(const string& s, int i, bool canChange,
                                   int mask, int k,
                                   unordered_map<long, int>& mem) {
    if (i == s.length())
      return 0;

    long key = static_cast<long>(i) << 27 | (canChange ? 1 : 0) << 26 | mask;
    if (const auto it = mem.find(key); it != mem.end())
      return it->second;

    // Initialize the result based on the current letter.
    int res = getRes(s, i, canChange, mask, 1 << (s[i] - 'a'), k, mem);

    // If allowed, explore the option to change the current letter.
    if (canChange)
      for (int j = 0; j < 26; ++j)
        res = max(res, getRes(s, i, false, mask, 1 << j, k, mem));

    return mem[key] = res;
  }

  int getRes(const string& s, int i, bool nextCanChange, unsigned mask,
             int newBit, int k, unordered_map<long, int>& mem) {
    const unsigned nextMask = mask | newBit;
    if (popcount(nextMask) > k)  // fresh start
      return 1 + maxPartitionsAfterOperations(s, i + 1, nextCanChange, newBit,
                                              k, mem);
    return maxPartitionsAfterOperations(s, i + 1, nextCanChange, nextMask, k,
                                        mem);
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class Solution {
  public int maxPartitionsAfterOperations(String s, int k) {
    Map<Long, Integer> mem = new HashMap<>();
    return maxPartitionsAfterOperations(s, 0, true, 0, k, mem) + 1;
  }

  // Returns the maximum number of partitions of s[i..n), where `canChange` is
  // true if we can still change a letter, and `mask` is the bitmask of the
  // letters we've seen.
  private int maxPartitionsAfterOperations(final String s, int i, boolean canChange, int mask,
                                           int k, Map<Long, Integer> mem) {
    if (i == s.length())
      return 0;

    Long key = (long) i << 27 | (canChange ? 1 : 0) << 26 | mask;
    if (mem.containsKey(key))
      return mem.get(key);

    // Initialize the result based on the current letter.
    int res = getRes(s, i, canChange, mask, 1 << (s.charAt(i) - 'a'), k, mem);

    // If allowed, explore the option to change the current letter.
    if (canChange)
      for (int j = 0; j < 26; ++j)
        res = Math.max(res, getRes(s, i, false, mask, 1 << j, k, mem));

    mem.put(key, res);
    return res;
  }

  private int getRes(final String s, int i, boolean nextCanChange, int mask, int newBit, int k,
                     Map<Long, Integer> mem) {
    final int nextMask = mask | newBit;
    if (Integer.bitCount(nextMask) > k) // fresh start
      return 1 + maxPartitionsAfterOperations(s, i + 1, nextCanChange, newBit, k, mem);
    return maxPartitionsAfterOperations(s, i + 1, nextCanChange, nextMask, k, mem);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution:
  def maxPartitionsAfterOperations(self, s: str, k: int) -> int:
    @functools.lru_cache(None)
    def dp(i: int, canChange: bool, mask: int) -> int:
      """
      Returns the maximum number of partitions of s[i..n), where `canChange` is
      True if we can still change a letter, and `mask` is the bitmask of the
      letters we've seen.
      """
      if i == len(s):
        return 0

      def getRes(newBit: int, nextCanChange: bool) -> int:
        nextMask = mask | newBit
        if nextMask.bit_count() > k:
          return 1 + dp(i + 1, nextCanChange, newBit)
        return dp(i + 1, nextCanChange, nextMask)

      # Initialize the result based on the current letter.
      res = getRes(1 << (ord(s[i]) - ord('a')), canChange)

      # If allowed, explore the option to change the current letter.
      if canChange:
        for j in range(26):
          res = max(res, getRes(1 << j, False))
      return res

    return dp(0, True, 0) + 1