Minimum sum of squares of character counts in a given string after removing k characters

The problem, simply

You get a string s of small letters and a number k. Count how many times each letter appears. The score is the sum of the squares of those counts: a letter that appears 3 times adds 9.

You must remove exactly k letters (any letters you like). Give back the smallest score you can end up with. If k is as big as the string, everything is removed and the score is 0.

A story to remember it

Picture it

Examples

Way 1: the simple way

Count the letters into 26 boxes. Then k times: look through all 26 boxes, find the biggest count, and lower it by one. Finally add up the squares.

1. count[c]++ for each letter.
2. Repeat k times: scan the 26 boxes for the biggest, subtract 1 (stop if all are 0).
3. Return the sum of count[c] * count[c].

class Solution {
    public int minValue(String s, int k) {
        int[] count = new int[26];
        for (char c : s.toCharArray()) count[c - 'a']++;
        for (int step = 0; step < k; step++) {
            int big = 0;
            for (int j = 1; j < 26; j++) {
                if (count[j] > count[big]) big = j;      // the tallest tower
            }
            if (count[big] == 0) break;                  // nothing left to remove
            count[big]--;
        }
        int score = 0;
        for (int c : count) score += c * c;
        return score;
    }
}

How long it takes: counting is O(n), and each of the k removals scans 26 boxes, so O(n + 26k). With a bigger alphabet of size m that becomes O(n + mk). Extra room: 26 boxes, O(1).

Way 2: the clever way

Why is "always the tallest" right? Lowering a count from f to f - 1 lowers the score by f * f - (f - 1) * (f - 1) = 2f - 1. That gain is biggest when f is biggest, and taking one bucket never makes any other tower's gain bigger later. So the greedy choice is always safe.

To find the tallest fast, keep the counts in a max heap, a priority queue that always hands out the biggest value first.

1. Count the letters and put every non zero count in a max heap.
2. Repeat k times: take the biggest count out, subtract 1, and put it back if it is still above 0.
3. Add up the squares of what is left in the heap.

import java.util.*;

class Solution {
    public int minValue(String s, int k) {
        int[] count = new int[26];
        for (char c : s.toCharArray()) count[c - 'a']++;
        PriorityQueue<Integer> heap = new PriorityQueue<>(Collections.reverseOrder());
        for (int c : count) if (c > 0) heap.offer(c);    // biggest count on top
        for (int step = 0; step < k && !heap.isEmpty(); step++) {
            int top = heap.poll() - 1;                    // shave the tallest tower
            if (top > 0) heap.offer(top);
        }
        int score = 0;
        while (!heap.isEmpty()) {
            int c = heap.poll();
            score += c * c;
        }
        return score;
    }
}

How long it takes: counting is O(n), and each removal is one heap take and one heap put, O(log 26), so O(n + k log 26). With an alphabet of size m it is O(n + k log m), which is the real win over Way 1. Extra room: at most 26 counts in the heap, O(1).

Another way: buckets of counts

Counts are small whole numbers, so skip the heap. Let howMany[f] be the number of letters whose count is f. Walk f from the largest count down. At level f, lower as many of those letters as you can (up to k) to f - 1. Then move to level f - 1, which now also holds the letters you just lowered.

class Solution {
    public int minValue(String s, int k) {
        int n = s.length();
        int[] count = new int[26];
        for (char c : s.toCharArray()) count[c - 'a']++;
        int[] howMany = new int[n + 1];                   // howMany[f] = letters with count f
        for (int c : count) if (c > 0) howMany[c]++;
        for (int f = n; f >= 1 && k > 0; f--) {
            int take = Math.min(k, howMany[f]);           // shave this many towers of height f
            howMany[f] -= take;
            howMany[f - 1] += take;
            k -= take;
        }
        int score = 0;
        for (int f = 1; f <= n; f++) score += howMany[f] * f * f;
        return score;
    }
}

This is O(n) time and O(n) extra room, and it works for any alphabet size.

Watch out for

• k bigger than the number of letters. Stop when the heap is empty, and return 0.
• Putting a count of 0 back into the heap. It does no harm to the score, but it wastes steps, so skip it.
• Overflow. With a string of a million copies of one letter, count * count is 10^12. Use long for the score if the limits are large.

What interviewers ask next

• Prove the greedy choice is right. Each removal gains 2f - 1, which grows with f. Swapping any removal to a shorter tower can only gain less, so no plan beats always shaving the tallest.
• Why does making counts equal help? The square function curves upward, so with a fixed total, the sum of squares is smallest when the parts are as equal as possible. Quant interviewers love this convexity argument.
• Do it when k is huge, like 10^9, on huge counts. Use the bucket idea in big jumps: lower the whole top level to the next level at once, and split the last partial step evenly.
• Rearrange a string so no two equal letters touch. Same max heap of counts, but you take the top two letters each round.

Remember this

Grown-up words: the towers are character frequencies, the lumpiness is the sum of squares, always taking the best move right now is a greedy algorithm, the line that serves the biggest first is a priority queue, built as a max heap, and the bucket trick is counting sort on frequencies.