LeetCodeWeekly238B: 1838. Frequency of the Most Frequent Element: 尺取り法


題意

  • $n$個の正の数字が与えられる。
  • 「あなたは好きな要素を$+1$してよい」これを最大で$k$回行える。(すべて別の要素を選んでもよい)
  • ある$x$を$y$個作りたいと思ったとき$y$の最大数はいくつか?

こう考えた

例題の通り、$[1,4,8,13]$としよう。まず、ソートする。今回の場合、すでにソートされている。
$+1$しかできないので、ある数字にそろえるとすると、その数より大きい数は小さくできない。つまり、$4$は$+1$を4回行えば$8$になれるが、$8$は操作しても$1$にも$4$にもなれない。

ここで左から順に、それまでの要素をいくつ揃えられるかを考える。例えば、13に合わせるとしよう。kが十分に大きければ$[13,13,13,13]$にできる($k$が本当に大きければ$[14,14,14,14]$などにもできるが、これはyの最大数に意味がない)。
必要な$k$が小さい方から考えると、$[1,4,13,13]$, $[1,13,13,13]$, $[13,13,13,13]$のように、13の一つ左の要素から何個目まで$k$で足りるかを考えるのが良さそうである。また、この時に必要な数は、$(一番右の数 \times 要素の数) - (その区間の和)$ である。わかりやすく考えるなら、$[4,8,13]$を$[13,13,13]$にするのには、$(13-4) + (13-8) + (13-13)$が必要であるが、これは先ほどの式に他ならない。

さて、今回のrは連続区間を求めるものなので、ある要素$l$, $r$番目を選んだ時に$(r-l+1)$を最大と言い換えられる。ここで、$r$を増加させていくと、$l$は必ず減算方向にはいかない。
このため、$l=0$の状態から$r$を引き延ばしていき、各段階で、$l$をどこまで$r$に近づければならないかを考えればよい。

つまり、for rで、lを引き寄せる尺取りをする。

実装

from typing import List, Tuple
from pprint import pprint
class Solution:
    def maxFrequency(self, nums: List[int], k: int) -> int:
        l = 0
        nums.sort()
        res = 0
        sum = 0
        for r in range(len(nums)):
            sum += nums[r]
            while True:
                targetnum = nums[r] * (r - l + 1)
                if k < (targetnum - sum):
                    sum -= nums[l]
                    l += 1
                else:
                    res = max(res, (r - l + 1))
                    break
        print(res)
        return res