py-leetcode/3347.py at main · uchitsa/py-leetcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class Solution:
    def maxFrequency(self, nums: List[int], k: int, numOperations: int) -> int:
        nums.sort()
        res = 0
        cnt = defaultdict(int)
        mds = set()

        def addmod(v):
            mds.add(v)
            if v-k >= nums[0]:
                mds.add(v-k)
            if v+k <= nums[-1]:
                mds.add(v+k)

        last = 0
        for i in range(len(nums)):
            if nums[i]!=nums[last]:
                cnt[nums[last]] = i-last
                res = max(res,i-last)
                addmod(nums[last])
                last = i

        cnt[nums[last]] = len(nums)-last
        res = max(res, len(nums)-last)
        addmod(nums[last])

        for m in sorted(mds):
            le = bisect.bisect_left(nums,m-k)
            ri = bisect.bisect_right(nums,m+k)-1
            if m in cnt:
                t = min(ri-le+1, cnt[m]+numOperations)
            else:
                t = min(ri-le+1, numOperations)
            res = max(res, t)
        return res