BinarySearch 0760 Prefix with Equivalent Frequencies

Problem statement

https://binarysearch.com/problems/Prefix-with-Equivalent-Frequencies/

Solution

Equal to Leetcode 1224. Maximum Equal Frequency.

Complexity

It is O(n) for time and space.

Code

class Solution:
    def solve(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