Problem statement

https://leetcode.com/problems/maximum-equal-frequency/

Solution

The idea is to keep Counter cnt of numbers as well as for each frequency keep how many numbers we have with this frequency. For example for [3, 3, 4, 4, 4, 5, 5, 6, 6, 7, 7, 7] we have frequencies [2, 3, 2, 2, 3] and hence we have freqs[2] = 3 and freqs[3] = 2. When we add new number, we:

  1. Check how many times we seen this number before, if we seen it before, decrease freqs[f] by 1.
  2. If we have empty freqs[f], that is we do not have elements with frequency f anymore, pop it, we do not want to have zero elements.
  3. Increase freqs[f + 1] by one, because now we have one more element with frequency f + 1. Also update cnt[num].

Now, we need to look at our freqs counter and check:

  1. If we have length more than 2, continue, it means that we have at least 3 different frequencies, we can not make them all equal.
  2. If we have length equal to 1, it means, that we have all element with the same frequency. In this case, two options are ok for us: either all frequencies equal to 1, like in [1, 2, 3, 4, 5] or we have only one value, like [1, 1, 1, 1, 1]. Other cases, like [1, 1, 2, 2, 3, 3] are not OK, because we need to remove exactly one element.
  3. If we have length equal to 2, then two options are OK for us: if we have pair (1, 1), (x, y), that is we have A, y_1, ..., y_1, ..., y_x, ... y_x (1 element of frequency 1 and x elements of frequency y) and we can remove the first one. Or we can have (1, y+1), (x, y) case, then we have one element of frequency y+1 and x elements of frequency y. Finally, we can have (x, y), (1, y+1) case, because we can not sure in which order we have elements in our t.

Complexity

Time complexity is O(n), space as well.

Code

class Solution:
    def maxEqualFreq(self, nums):
        cnt, freqs, ans = Counter(), Counter(), 1

        for i, num in enumerate(nums):
            f = cnt[num]
            if f > 0: freqs[f] -= 1
            if f in freqs and freqs[f] == 0: freqs.pop(f)
            freqs[f + 1] += 1
            cnt[num] += 1
            
            if len(freqs) > 2: continue
            t = list(freqs.items())
             
            if len(t) == 1 and 1 in t[0]:
                ans = max(ans, i)
                
            if len(t) == 2:
                (a, b), (c, d) = t
                if (1, 1) in t or b*(a-c) == 1 or d*(c-a) == 1:
                    ans = max(ans, i)
             
        return ans + 1