Maximize the number of partitions after operations

impl Solution {
  pub fn max_partitions_after_operations(s: String, k: i32) -> i32 {
    let s: Vec<u8> = s.bytes().map(|b| b - b'a').collect();
    let n = s.len();
    let k = k as usize;

    let ne = Self::next_ends(&s, n, k);
    let ne1 = if k > 0 { Self::next_ends(&s, n, k - 1) } else { vec![n; n + 1] };

    // first_force[i]: first p in [i, ne[i]) where |distinct(s[i..p-1])| = k and s[p] in window
    // also stores mask_before = distinct(s[i..p-1]) as u32
    let mut first_force = vec![None::<(usize, u32)>; n + 1];
    for i in 0..n {
      let end = ne[i];
      let mut mask: u32 = 0;
      let mut cnt = 0usize;
      for p in i..end {
        if cnt == k && (mask >> s[p]) & 1 == 1 {
          first_force[i] = Some((p, mask));
          break;
        }
        if (mask >> s[p]) & 1 == 0 { cnt += 1; mask |= 1 << s[p]; }
      }
    }

    // Precompute next_pos[i][c] = next position >= i where char c appears.
    // mask of s[a..b-1] = bitset of chars c where next_pos[a][c] < b  → O(26) per query
    let mut next_pos = vec![[n; 26usize]; n + 1];
    for i in (0..n).rev() {
      next_pos[i] = next_pos[i + 1];
      next_pos[i][s[i] as usize] = i;
    }

    // mask_window[p] = distinct char mask of s[p+1..ne1[p+1]-1]
    let mask_window: Vec<u32> = (0..n).map(|p| {
      let start = p + 1;
      if start >= n { return 0u32; }
      let end = ne1[start];
      let mut m = 0u32;
      for c in 0..26usize {
        if next_pos[start][c] < end { m |= 1 << c; }
      }
      m
    }).collect();

    let mut dp = vec![[0i32; 2]; n + 1];
    for i in (0..n).rev() {
      dp[i][1] = 1 + dp[ne[i]][1];
      let opt1 = 1 + dp[ne[i]][0];
      let opt2 = match first_force[i] {
        // Guard: mask_before has k distinct chars. If k=26 every character is
        // already in the window, so there is no c ∉ mask_before to change s[p]
        // to — the change would be invalid and can produce no extra partition.
        Some((p, mask_before)) if mask_before != (1u32 << 26) - 1 => {
          // We change s[p] to c ∉ distinct(s[i..p-1]).
          // To achieve the shortest next partition (end at ne1[p+1]):
          //   need c ∉ mask_before ∪ mask_window[p] ∪ {s[ne1[p+1]]}
          // Otherwise we must use ne[p+1] (c is already in the k-1 window).
          let next = if p + 1 < n {
            let mw = mask_window[p];
            let barrier = if ne1[p + 1] < n { 1u32 << s[ne1[p + 1]] } else { 0 };
            let blocked = mask_before | mw | barrier;
            if blocked != (1u32 << 26) - 1 {
              ne1[p + 1]   // can pick c outside all windows → shorter partition
            } else {
              ne[p + 1]    // forced into longer partition
            }
          } else {
            n
          };
          2 + dp[next][1]
        }
        _ => 0,
      };
      // opt3: force split at ne1[i] by changing a char in [i..ne1[i]-1] to a new distinct,
      // making [i..ne1[i]-1] have k distinct and starting next partition at ne1[i] with change used.
      // Requires at least k chars in [i..ne1[i]-1] so the change adds a new distinct without losing an existing one.
      // Also requires k < 26: the window has k-1 distinct, leaving 27-k free chars; one of them is
      // s[ne1[i]] itself. We need at least one OTHER free char to change to — if k=26 the only free
      // char is s[ne1[i]], so after the change s[ne1[i]] is already in the window and no split occurs.
      let opt3 = if ne1[i] >= i + k && ne1[i] < ne[i] && k < 26 { 1 + dp[ne1[i]][1] } else { 0 };
      // opt4: change a UNIQUE char s[j] in [i..ne1[i]-1] (first occurrence = j) to a new distinct c'.
      // After the change, s[j]'s original value is absent from the window. When it reappears at
      // next_occ > ne1[i], the window [i..next_occ-1] has exactly k distinct (proven: original k
      // minus removed char plus c') and s[next_occ] adds the (k+1)-th → forced split at next_occ.
      // Requires ne1[i] < next_occ < ne[i] for the split to occur within the original partition.
      let opt4 = if ne1[i] < ne[i] && k < 26 {
        let mut best = 0i32;
        for c in 0..26usize {
          let first_occ = next_pos[i][c];
          // When i > 0 and first_occ == i, changing s[i] (the overflow char that ends
          // the previous partition) to a char already in that window would extend the
          // previous partition past i, invalidating the boundary assumed by this DP.
          // Skip that case; when i == 0 there is no previous partition to worry about.
          if first_occ < ne1[i] && (first_occ > i || i == 0) {
            let next_occ = next_pos[first_occ + 1][c];
            if next_occ > ne1[i] && next_occ < ne[i] {
              let cand = 1 + dp[next_occ][1];
              if cand > best { best = cand; }
            }
          }
        }
        best
      } else { 0 };
      dp[i][0] = opt1.max(opt2).max(opt3).max(opt4);
    }
    dp[0][0]
  }

  fn next_ends(s: &[u8], n: usize, k: usize) -> Vec<usize> {
    if k == 0 { return (0..=n).collect(); }
    let mut ne = vec![n; n + 1];
    let mut freq = [0u32; 26];
    let mut distinct = 0usize;
    let mut right = 0usize;
    for left in 0..n {
      loop {
        if right >= n { ne[left] = n; break; }
        let c = s[right] as usize;
        let new_d = if freq[c] == 0 { distinct + 1 } else { distinct };
        if new_d > k { ne[left] = right; break; }
        if freq[c] == 0 { distinct += 1; }
        freq[c] += 1;
        right += 1;
      }
      let c = s[left] as usize;
      if freq[c] > 0 {
        freq[c] -= 1;
        if freq[c] == 0 { distinct -= 1; }
      }
    }
    ne
  }
}