Count runs with exactly k distinct values

from collections import defaultdict

def count_exactly_k(nums, k):
    def at_most(k):
        count = defaultdict(int)
        total = lo = 0
        for hi, x in enumerate(nums):
            count[x] += 1
            while len(count) > k:            # too many values: shrink
                count[nums[lo]] -= 1
                if count[nums[lo]] == 0:
                    del count[nums[lo]]
                lo += 1
            total += hi - lo + 1             # every start from lo to hi has at most k
        return total

    return at_most(k) - at_most(k - 1)

#include <vector>

static long long atMost(const std::vector<int>& nums, int k) {
    std::vector<int> count(nums.size() + 1, 0);      // values are 1..n
    long long total = 0;
    int kinds = 0, lo = 0;
    for (int hi = 0; hi < (int)nums.size(); ++hi) {
        if (count[nums[hi]]++ == 0) ++kinds;
        while (kinds > k) if (--count[nums[lo++]] == 0) --kinds;   // too many values: shrink
        total += hi - lo + 1;                        // every start from lo to hi has at most k
    }
    return total;
}

long long solve(const std::vector<int>& nums, int k) {
    return atMost(nums, k) - atMost(nums, k - 1);
}

class Solution {
    private long atMost(int[] nums, int k) {
        int[] count = new int[nums.length + 1];         // values are 1..n
        long total = 0;
        int kinds = 0, lo = 0;
        for (int hi = 0; hi < nums.length; hi++) {
            if (count[nums[hi]]++ == 0) kinds++;
            while (kinds > k) if (--count[nums[lo++]] == 0) kinds--;   // too many values: shrink
            total += hi - lo + 1;
        }
        return total;
    }

    public long countExactlyK(int[] nums, int k) {
        return atMost(nums, k) - atMost(nums, k - 1);
    }
}

function countExactlyK(nums, k) {
  const atMost = k => {
    const count = new Map();
    let total = 0, lo = 0;
    for (let hi = 0; hi < nums.length; hi++) {
      count.set(nums[hi], (count.get(nums[hi]) || 0) + 1);
      while (count.size > k) {               // too many values: shrink
        const n = count.get(nums[lo]) - 1;
        if (n === 0) count.delete(nums[lo]); else count.set(nums[lo], n);
        lo++;
      }
      total += hi - lo + 1;                  // every start from lo to hi has at most k
    }
    return total;
  };
  return atMost(k) - atMost(k - 1);
}

fn at_most(nums: &[i64], k: usize) -> i64 {
    let mut count = vec![0usize; nums.len() + 1];      // values are 1..n
    let (mut kinds, mut lo, mut total) = (0usize, 0usize, 0i64);
    for hi in 0..nums.len() {
        let x = nums[hi] as usize;
        if count[x] == 0 { kinds += 1; }
        count[x] += 1;
        while kinds > k {                              // too many values: shrink
            let y = nums[lo] as usize;
            count[y] -= 1;
            if count[y] == 0 { kinds -= 1; }
            lo += 1;
        }
        total += (hi + 1 - lo) as i64;
    }
    total
}

fn solve(nums: &[i64], k: usize) -> i64 {
    at_most(nums, k) - at_most(nums, k - 1)
}

import qualified Data.Map.Strict as M

-- Exactly k = at most k − at most (k − 1); each pass is a window that counts its starts.
solve :: [Int] -> Int -> Int
solve nums k = atMost k - atMost (k - 1)
  where
    atMost kk = total
      where
        (_, _, _, total) = foldl step (M.empty, nums, 0, 0) nums
        step (m, win, len, t) x =
          let (m', win', len') = shrink (M.insertWith (+) x 1 m) win (len + 1)
          in (m', win', len', t + len')
        shrink m win len
          | M.size m > kk, (y:ys) <- win =
              shrink (M.update (\c -> if c == 1 then Nothing else Just (c - 1)) y m) ys (len - 1)
          | otherwise = (m, win, len)