Maximize the Number of Partitions After Operations
Source: leetcode
Difficulty: Hard
Tags: string, dynamic-programming, bit-manipulation, bitmask

You are given a string `s` and an integer `k`.


First, you are allowed to change **at most** **one** index in `s` to another lowercase English letter.


After that, do the following partitioning operation until `s` is **empty**:


- Choose the **longest** **prefix** of `s` containing at most `k` **distinct** characters.
- **Delete** the prefix from `s` and increase the number of partitions by one. The remaining characters (if any) in `s` maintain their initial order.


Return an integer denoting the **maximum** number of resulting partitions after the operations by optimally choosing at most one index to change.


 


**Example 1:**



**Input:** s \= "accca", k \= 2


**Output:** 3


**Explanation:**


The optimal way is to change `s[2]` to something other than a and c, for example, b. then it becomes `"acbca"`.


Then we perform the operations:


1. The longest prefix containing at most 2 distinct characters is `"ac"`, we remove it and `s` becomes `"bca"`.
2. Now The longest prefix containing at most 2 distinct characters is `"bc"`, so we remove it and `s` becomes `"a"`.
3. Finally, we remove `"a"` and `s` becomes empty, so the procedure ends.


Doing the operations, the string is divided into 3 partitions, so the answer is 3\.



**Example 2:**



**Input:** s \= "aabaab", k \= 3


**Output:** 1


**Explanation:**


Initially `s` contains 2 distinct characters, so whichever character we change, it will contain at most 3 distinct characters, so the longest prefix with at most 3 distinct characters would always be all of it, therefore the answer is 1\.



**Example 3:**



**Input:** s \= "xxyz", k \= 1


**Output:** 4


**Explanation:**


The optimal way is to change `s[0]` or `s[1]` to something other than characters in `s`, for example, to change `s[0]` to `w`.


Then `s` becomes `"wxyz"`, which consists of 4 distinct characters, so as `k` is 1, it will divide into 4 partitions.



 


**Constraints:**


- `1 <= s.length <= 104`
- `s` consists only of lowercase English letters.
- `1 <= k <= 26`




C++ solution:
// Time:  O(n)
// Space: O(n)

// prefix sum, greedy
class Solution {
public:
    int maxPartitionsAfterOperations(string s, int k) {
        vector<int> left(size(s) + 1);
        vector<int> left_mask(size(s) + 1);
        for (int i = 0, cnt = 0, mask = 0; i < size(s); ++i) {
            mask |= 1 << (s[i] - 'a');
            if (__builtin_popcount(mask) > k) {
                ++cnt;
                mask = 1 << (s[i] - 'a');
            }
            left[i + 1] = cnt;
            left_mask[i + 1] = mask;
        }
        vector<int> right(size(s) + 1);
        vector<int> right_mask(size(s) + 1);
        for (int i = size(s) - 1, cnt = 0, mask = 0; i >= 0; --i) {
            mask |= 1 << (s[i] - 'a');
            if (__builtin_popcount(mask) > k) {
                ++cnt;
                mask = 1 << (s[i] - 'a');
            }
            right[i] = cnt;
            right_mask[i] = mask;
        }

        int result = 0;
        for (int i = 0; i < size(s); ++i) {
            int curr = left[i] + right[i + 1];
            const int mask = left_mask[i] | right_mask[i + 1];
            if (__builtin_popcount(left_mask[i]) == k && __builtin_popcount(right_mask[i + 1]) == k && __builtin_popcount(mask) != 26) {
                curr += 3;
            } else if (__builtin_popcount(mask) + (__builtin_popcount(mask) != 26 ? 1 : 0) > k) {  // test case: s = "abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz", k = 26
                curr += 2;
            } else {
                curr += 1;
            }
            result = max(result, curr);
        }
        return result;
    }
};

Go solution:
func maxPartitionsAfterOperations(s string, k int) int {
    n := len(s)
    if n == 0 {
        return 0
    }

    left := make([]int, n+1)
    leftMask := make([]int, n+1)
    currentMask := 0
    currentPartitions := 0
    for i := 0; i < n; i++ {
        currentMask |= 1 << (s[i] - 'a')
        if bits.OnesCount(uint(currentMask)) > k {
            currentPartitions++
            currentMask = 1 << (s[i] - 'a')
        }
        left[i+1] = currentPartitions
        leftMask[i+1] = currentMask
    }

    right := make([]int, n+1)
    rightMask := make([]int, n+1)
    currentMask = 0
    currentPartitions = 0
    for i := n - 1; i >= 0; i-- {
        currentMask |= 1 << (s[i] - 'a')
        if bits.OnesCount(uint(currentMask)) > k {
            currentPartitions++
            currentMask = 1 << (s[i] - 'a')
        }
        right[i] = currentPartitions
        rightMask[i] = currentMask
    }

    maxResult := 0
    for i := 0; i < n; i++ {
        curr := left[i] + right[i+1]
        mask := leftMask[i] | rightMask[i+1]
        leftCount := bits.OnesCount(uint(leftMask[i]))
        rightCount := bits.OnesCount(uint(rightMask[i+1]))
        maskCount := bits.OnesCount(uint(mask))

        var temp int
        if maskCount != 26 {
            temp = maskCount + 1
        } else {
            temp = maskCount
        }

        if leftCount == k && rightCount == k && maskCount != 26 {
            curr += 3
        } else if temp > k {
            curr += 2
        } else {
            curr += 1
        }

        if curr > maxResult {
            maxResult = curr
        }
    }

    return maxResult
}

Java solution:
class Solution {

    static final int ALPHABET_SIZE = 'z' - 'a' + 1;

    public int maxPartitionsAfterOperations(String s, int k) {
        if (k == ALPHABET_SIZE) {
            return 1;
        }
        int n = s.length();

        int[] ansr = new int[n];
        int[] usedr = new int[n];
        int used = 0;
        int cntUsed = 0;
        int ans = 1;
        for (int i = n - 1; i >= 0; --i) {
            int ch = s.charAt(i) - 'a';
            if ((used & (1 << ch)) == 0) {
                if (cntUsed == k) {
                    cntUsed = 0;
                    used = 0;
                    ans++;
                }
                used |= (1 << ch);
                cntUsed++;
            }
            ansr[i] = ans;
            usedr[i] = used;
        }

        int ansl = 0;
        ans = ansr[0];
        for (int l = 0; l < n;) {
            used = 0;
            cntUsed = 0;
            int usedBeforeLast = 0;
            int usedTwiceBeforeLast = 0;
            int last = -1;
            int r = l;
            while (r < n) {
                int ch = s.charAt(r) - 'a';
                if ((used & (1 << ch)) == 0) {
                    if (cntUsed == k) {
                        break;
                    }
                    usedBeforeLast = used;
                    last = r;
                    used |= (1 << ch);
                    cntUsed++;
                } else if (cntUsed < k) {
                    usedTwiceBeforeLast |= (1 << ch);
                }
                r++;
            }

            if (cntUsed == k) {
                if (last - l > Integer.bitCount(usedBeforeLast)) {
                    // we can change one letter in s[l..last), so new partition will start at last
                    ans = Math.max(ans, ansl + 1 + ansr[last]);
                }
                if (last + 1 < r) {
                    // we can change s[last+1] to a letter to maximize the number of partitions
                    // s[l..last] is one partition, so new partition starts from s[last + 1]
                    if (last + 2 >= n) {
                        ans = Math.max(ans, ansl + 1 + 1);
                    } else {
                        if (Integer.bitCount(usedr[last + 2]) == k) {
                            int canUse = ((1 << ALPHABET_SIZE) - 1) & ~used & ~usedr[last + 2];
                            if (canUse > 0) {
                                ans = Math.max(ans, ansl + 1 + 1 + ansr[last + 2]);
                            } else {
                                ans = Math.max(ans, ansl + 1 + ansr[last + 2]);
                            }
                            int l1 = s.charAt(last + 1) - 'a';
                            if ((usedTwiceBeforeLast & (1 << l1)) == 0) {
                                ans = Math.max(ans, ansl + 1 + ansr[last + 1]);
                            }
                        } else {
                            ans = Math.max(ans, ansl + 1 + ansr[last + 2]);
                        }
                    }
                }
            }

            l = r;
            ansl++;
        }

        return ans;
    }
}

JavaScript solution:
let FULL_CHAR_STATE = 67108863; // 2^26 - 1
let counters = new Uint16Array(26);

/**
 * The largest index r such that [i, r) has at most k distinct chars
 */
let longestKs = new Uint16Array(10002);
/**
 * The largest index r such that [i, r) has at most k-1 distinct chars
 */
let longestK1s = new Uint16Array(10002);

let encodedKs = new Uint32Array(10002);
let dp = new Uint16Array(10002);

function canAdd(code1, code2) {
  return (code1 | code2) != FULL_CHAR_STATE;
}

/**
 * @param {string} s
 * @param {number} k
 * @return {number}
 */
var maxPartitionsAfterOperations = function (s, k) {
  if (k == 26) return 1;

  let n = s.length;
  if (k == 1) {
    let res = 0;
    let addition = 0;
    for (let i = 0; i < n; ++i) {
      if (s[i] != s[i - 1]) ++res;
      else {
        addition = Math.max(addition, 1);
        if (s[i] == s[i + 1]) addition = 2;
      }
    }

    return res + addition;
  }

  let cnt = 0;
  let code = 0;
  let r = -1;

  counters.fill(0);
  for (let l = 0; l < n; ++l) {
    if (l) {
      let ch = s.charCodeAt(l - 1) - 97;
      --counters[ch];
      if (!counters[ch]) {
        --cnt;
      }
    }
    while (r < n && cnt < k) {
      ++r;
      if (r == n) break;
      let ch = s.charCodeAt(r) - 97;
      ++counters[ch];
      if (counters[ch] == 1) {
        ++cnt;
      }
    }

    longestK1s[l] = r;
  }

  counters.fill(0);
  cnt = 0;
  r = -1;
  for (let l = 0; l < n; ++l) {
    if (l) {
      let ch = s.charCodeAt(l - 1) - 97;
      --counters[ch];
      if (!counters[ch]) {
        --cnt;
        code &= ~(1 << ch);
      }
    }
    while (r < n && cnt <= k) {
      ++r;
      if (r == n) break;
      let ch = s.charCodeAt(r) - 97;
      if (!counters[ch]) {
        ++cnt;
        code |= 1 << ch;
      }
      ++counters[ch];
    }
    longestKs[l] = r;
    encodedKs[l] = code;
    if (r < n) {
      encodedKs[l] ^= 1 << (s.charCodeAt(r) - 97);
    }
  }
  encodedKs[n] = 0;
  encodedKs[n + 1] = 0;
  encodedKs[n + 2] = 0;
  longestKs[n] = n;
  longestKs[n + 1] = n;
  longestKs[n + 2] = n;
  longestK1s[n] = n;
  longestK1s[n + 1] = n;
  longestK1s[n + 2] = n;

  dp.fill(0, 0, n + 3);

  for (let i = n - 1; i >= 0; --i) {
    dp[i] = dp[longestKs[i]] + 1;
  }

  let res = dp[0];
  let currStart = 0;
  let currEnd = longestKs[0];
  let prevCode = 0;
  let currCode = encodedKs[0];
  let cntPre = 1;

  while (true) {
    let lk = currEnd - currStart;
    if (lk > k) {
      let lk1 = longestK1s[currStart] - currStart;

      // if there are k consecutive different chars
      if (lk1 + 1 == k) {
        // changing the char at currStart + k
        if (canAdd(currCode, encodedKs[currStart + k + 1])) {
          res = Math.max(res, cntPre + 1 + dp[longestK1s[currStart + k + 1]]);
        }

        // changing the first occurrence of the first duplicated char
        if (
          s[currStart] != s[currStart + k] ||
          canAdd(prevCode, encodedKs[currStart + 1])
        ) {
          res = Math.max(res, cntPre + 1 + dp[longestKs[currStart + k]]);
        }
      } else {
        // changing the second occurrence of the first duplicated char
        res = Math.max(res, cntPre + dp[longestK1s[currStart]]);

        // changing the first char x such that the subarray [currStart, x) has K distinct chars
        if (longestK1s[currStart] < n && canAdd(currCode, encodedKs[longestKs[currStart] + 2])) {
          res = Math.max(
            res,
            cntPre + 1 + dp[longestK1s[longestK1s[currStart] + 2]]
          );
        }
      }
    }

    if (currEnd == n) break;
    currStart = currEnd;
    currEnd = longestKs[currEnd];

    prevCode = currCode;
    currCode = encodedKs[currStart];
    ++cntPre;
  }

  return res;
};

Python solution:
class Solution:
    def maxPartitionsAfterOperations(self, s: str, k: int) -> int:
        if k == 26 or k > len(set(s)):  return 1

        if k == 1:
            mx, cur, cnt, prev = 0, 0, -1, ''
            for c in s:
                if c == prev:
                    cur += 1
                else:
                    mx = max(mx, cur)
                    cur, prev = 1, c
                    cnt += 1
            mx = max(mx, cur)
            return cnt + min(mx, 3)
        
        n, s = len(s), [1 << ord(c)-97 for c in s]

        suff1, suff2, setMap, inds = [0]*n, [0]*n, [0]*n, {1<<i: n for i in range(26)}

        set1, cnt1, par1, ptr1 = 0, 0, 1, n-1
        set2, cnt2, par2, ptr2 = 0, 0, 1, n-1

        for i in range(n-1,-1,-1):

            suff2[i], setMap[i] = par2, set1

            if not set2 & s[i]:
                if cnt2 == k-1:
                    while ptr2 > inds[s[ptr2]]:    ptr2 -= 1
                    par2 = suff1[ptr2] + 1
                    set2 ^= s[ptr2]
                    ptr2 -= 1
                    cnt2 -= 1
                set2 |= s[i]
                cnt2 += 1

            if not set1 & s[i]:
                if cnt1 == k:
                    while ptr1 > inds[s[ptr1]]:    ptr1 -= 1
                    par1 = suff1[ptr1] + 1
                    set1 ^= s[ptr1]
                    ptr1 -= 1
                    cnt1 -= 1
                set1 |= s[i]
                cnt1 += 1
            
            suff1[i], inds[s[i]] = par1, i

        # `res` starts from partitions without replacement
        res = suff1[0]
        mode1, mode2 = False, False
        set0, cnt0, par0, mask = 0, 0, 1, (1<<26)-1

        # mode1: Replace a redundant char before the kth char,
        #       new partition starts from the kth char;
        # mode2: Replace a redundant char after the kth char,
        #       new partition starts from the first redundant char
        #       after the kth char;
        #
        # e.g.:  (s,k) = ("...abbcac...", 3);
        #           mode1: replace the second 'b' (before the first 'c')
        #                   new partition starts from the first 'c'
        #           mode2: replace the second 'a' (after the first 'c')
        #                   new partition starts from the second 'a'

        for i in range(n):

            if not set0 & s[i]:

                if cnt0 == k-1:
                    # found the first kth char
                    if mode1:
                        res = max(res, par0 + suff1[i])
                    mode2 = True

                elif cnt0 == k:
                    # new partition
                    mode1, mode2 = False, False
                    set0, cnt0 = 0, 0
                    par0 += 1

                set0 |= s[i]
                cnt0 += 1

            elif mode2:
                # found the first redundant char after the kth char
                if set0 | setMap[i] < mask:
                    # Replace it with a distinct char
                    res = max(res, par0 + suff2[i])
                else:
                    # No distinct char available
                    res = max(res, par0 + suff1[i])
                mode2 = False

            elif not mode1:
                # found the first redundant char in the current partition
                mode1 = True
            
        return res

Ruby solution:
def max_partitions_after_operations(s, k)
    n = s.size
    return 1 if k == 26

    # Precompute left partitions and masks
    left = Array.new(n + 1, 0)
    left_mask = Array.new(n + 1, 0)
    current_mask = 0
    current_partitions = 0

    (0...n).each do |i|
        char = s[i]
        mask_bit = 1 << (char.ord - 'a'.ord)
        new_mask = current_mask | mask_bit
        if new_mask.to_s(2).count('1') > k
            current_partitions += 1
            current_mask = mask_bit
        else
            current_mask = new_mask
        end
        left[i + 1] = current_partitions
        left_mask[i + 1] = current_mask
    end

    # Precompute right partitions and masks
    right = Array.new(n + 1, 0)
    right_mask = Array.new(n + 1, 0)
    current_mask = 0
    current_partitions = 0

    (n - 1).downto(0) do |i|
        char = s[i]
        mask_bit = 1 << (char.ord - 'a'.ord)
        new_mask = current_mask | mask_bit
        if new_mask.to_s(2).count('1') > k
            current_partitions += 1
            current_mask = mask_bit
        else
            current_mask = new_mask
        end
        right[i] = current_partitions
        right_mask[i] = current_mask
    end

    result = 0

    (0...n).each do |i|
        curr = left[i] + right[i + 1]
        mask = left_mask[i] | right_mask[i + 1]

        count_left = left_mask[i].to_s(2).count('1')
        count_right = right_mask[i + 1].to_s(2).count('1')
        count_combined = mask.to_s(2).count('1')

        if count_left == k && count_right == k && count_combined != 26
            curr += 3
        elsif (count_combined + (count_combined != 26 ? 1 : 0)) > k
            curr += 2
        else
            curr += 1
        end

        result = [result, curr].max
    end

    # Consider the original scenario (no changes)
    original_partitions = left[n] + 1
    [result, original_partitions].max
end