Maximum Equal Frequency

hard array hash table frequency

Problem

Given an array nums of positive integers, return the longest possible length of a prefix of nums such that it is possible to remove exactly one element from this prefix so that every number that remains appears the same number of times.

If after removing one element there are no remaining elements, it is still considered that every appearing number has the same number of occurrences (0).

Inputnums = [2, 2, 1, 1, 5, 3, 3, 5]

Output7

For the prefix [2,2,1,1,5,3,3] of length 7, removing one occurrence of 5 leaves 2,1,3 each appearing exactly twice.

nums (comma-separated)

def max_equal_freq(nums):
    cnt = {}        # value -> occurrences
    freq = {}       # occurrences -> how many values have it
    ans = 0
    for i, x in enumerate(nums):
        c = cnt.get(x, 0)
        if c:
            freq[c] -= 1
        cnt[x] = c + 1
        freq[c + 1] = freq.get(c + 1, 0) + 1
        L = i + 1
        keys = [k for k, v in freq.items() if v > 0]
        lo, hi = min(keys), max(keys)
        if len(keys) == 1:
            # all equal freq: valid if it is 1 (drop any), or a single value
            ok = lo == 1 or freq[lo] == 1
        elif len(keys) == 2:
            # top freq one above the rest with a single owner, or one
            # lone value of freq 1 sitting on a uniform block
            ok = (hi - lo == 1 and freq[hi] == 1) or (lo == 1 and freq[lo] == 1)
        else:  # three or more distinct freqs can never be fixed by one removal
            ok = False
        if ok:
            ans = L
    return ans

function maxEqualFreq(nums) {
  const cnt = new Map();   // value -> occurrences
  const freq = new Map();  // occurrences -> how many values have it
  let ans = 0;
  for (let i = 0; i < nums.length; i++) {
    const x = nums[i];
    const c = cnt.get(x) || 0;
    if (c) freq.set(c, freq.get(c) - 1);
    cnt.set(x, c + 1);
    freq.set(c + 1, (freq.get(c + 1) || 0) + 1);
    const L = i + 1;
    const keys = [];
    for (const k of freq.keys()) if (freq.get(k) > 0) keys.push(k);
    const lo = Math.min(...keys), hi = Math.max(...keys);
    let ok;
    if (keys.length === 1) {
      ok = lo === 1 || freq.get(lo) === 1;
    } else if (keys.length === 2) {
      ok = (hi - lo === 1 && freq.get(hi) === 1) || (lo === 1 && freq.get(lo) === 1);
    } else { // three or more distinct freqs can never be fixed by one removal
      ok = false;
    }
    if (ok) ans = L;
  }
  return ans;
}

class Solution {
    public int maxEqualFreq(int[] nums) {
        Map<Integer, Integer> cnt = new HashMap<>();
        Map<Integer, Integer> freq = new HashMap<>();
        int ans = 0;
        for (int i = 0; i < nums.length; i++) {
            int x = nums[i];
            int c = cnt.getOrDefault(x, 0);
            if (c > 0) freq.put(c, freq.get(c) - 1);
            cnt.put(x, c + 1);
            freq.merge(c + 1, 1, Integer::sum);
            int L = i + 1, lo = Integer.MAX_VALUE, hi = 0, distinct = 0;
            for (int k : freq.keySet()) {
                if (freq.get(k) > 0) { distinct++; lo = Math.min(lo, k); hi = Math.max(hi, k); }
            }
            boolean ok;
            if (distinct == 1) {
                ok = lo == 1 || freq.get(lo) == 1;
            } else if (distinct == 2) {
                ok = (hi - lo == 1 && freq.get(hi) == 1) || (lo == 1 && freq.get(lo) == 1);
            } else { // three or more distinct freqs can never be fixed by one removal
                ok = false;
            }
            if (ok) ans = L;
        }
        return ans;
    }
}

int maxEqualFreq(vector<int>& nums) {
    unordered_map<int, int> cnt, freq;
    int ans = 0;
    for (int i = 0; i < (int)nums.size(); i++) {
        int x = nums[i];
        int c = cnt[x];
        if (c) freq[c]--;
        cnt[x] = c + 1;
        freq[c + 1]++;
        int L = i + 1, lo = INT_MAX, hi = 0, distinct = 0;
        for (auto& p : freq)
            if (p.second > 0) { distinct++; lo = min(lo, p.first); hi = max(hi, p.first); }
        bool ok;
        if (distinct == 1)
            ok = lo == 1 || freq[lo] == 1;
        else if (distinct == 2) // exactly two distinct freqs
            ok = (hi - lo == 1 && freq[hi] == 1) || (lo == 1 && freq[lo] == 1);
        else // three or more distinct freqs can never be fixed by one removal
            ok = false;
        if (ok) ans = L;
    }
    return ans;
}

Explanation

We scan prefixes from left to right and, at each length, ask: can removing exactly one element make every remaining value appear the same number of times? To check this instantly we keep two maps as we go.

cnt maps each value to how many times it has appeared so far. freq maps a count to how many distinct values currently have that count — a "count of the counts". When we add a value x, we decrement its old count's bucket in freq and increment the new one.

A prefix is valid only in a few tidy shapes. If there is just one live frequency, it works when that frequency is 1 (drop any single element) or only one value has it (drop that whole value). If there are exactly two frequencies, it works when the top is one higher than the rest and a single value owns it (drop one copy), or one lone value has frequency 1 sitting on an otherwise uniform block (drop that value).

Example: nums = [2,2,1,1,5,3,3,5]. At the prefix [2,2,1,1,5,3,3] of length 7, the counts are 2,1,3 each twice and 5 once; removing that single 5 leaves everyone at twice — valid, so the answer is 7.

Each step does constant work over a tiny set of frequencies, so the whole scan is linear.

Time: O(n) Space: O(n)