Maximum Equal Frequency LeetCode solution in Python

1
class Solution:
2
    def maxEqualFreq(self, nums: List[int]) -> int:
3
        ans = 0
4
        n = len(nums)
5
        countToFreq = defaultdict(int)
6
        # key = count value = Freq ex 2 occured 3 times  in nums so 2 : 3
7
        freqToCount = defaultdict(int)
8
        # key = freq value = count ex 2 numbers occured 3 times in nums so 2 : 3
9

10
        for i, val in enumerate(nums):
11

12
            x = countToFreq[val] + 1
13
            freqToCount[x - 1] -= 1
14
            if freqToCount[x - 1] <= 0:
15
                freqToCount.pop(x - 1)
16
            freqToCount[x] += 1
17
            countToFreq[val] = x
18

19
            # if a single item is repeated for i + 1 times like [1,1,1]
20
            if countToFreq[val] == i + 1:
21
                ans = i + 1
22

23
            # if all items are having same frequency like [2,2,1,1,3,3]
24
            elif (i < n - 1 and len(freqToCount) == 1) or (
25
                len(freqToCount) == 1 and max(freqToCount.keys()) == 1
26
            ):
27
                ans = i + 1
28

29
            # if all items have same frequency except one having 1 freq like [2,2,3,3,1]
30
            elif len(freqToCount) == 2 and 1 in freqToCount and freqToCount[1] == 1:
31
                ans = i + 1
32

33
            # if all items have same frequenct except one having freq common + 1 like[1,1,2,2,3,3,3]
34
            elif len(freqToCount) == 2:
35
                keys, values = [], []
36
                for j in freqToCount:
37
                    keys.append(j), values.append(freqToCount[j])
38
                if (keys[0] == 1 + keys[1] and values[0] == 1) or (
39
                    keys[1] == 1 + keys[0] and values[1] == 1
40
                ):
41
                    ans = i + 1
42
        return ans