410小R的数组构造挑战

题解

问题描述

小 R 有一个长度为 n 的数组，数组相邻元素的差值最多为 1，即 $|a_i − a_{i+1}| ≤ 1$，且数组中的元素都是正整数，即 $a_i ≥ 1$。现在已知数组的长度为 n，数组的和为 m，小 R 想知道所有符合条件的数组中，第 $p$ 个位置的元素 $a_p$ 的最大值是多少。

约束条件：

• 1 ≤ p ≤ n ≤ m

解题思路

1. 二分查找：

   • 使用二分查找在范围 [1, m] 内寻找 a_p 的最大可能值。
   • l 表示当前范围的下界，r 表示上界。

2. 辅助函数 calc(i, v)：

   • 计算从位置 i 开始，向两边扩展时，每个位置的值如何分配，使得相邻元素差值不超过 1。
   • 如果 v > i，则左侧的元素从 v - i + 1 递增到 v，形成等差数列。
   • 如果 v ≤ i，则前 v 个元素递增，剩下的元素为 1。
   • 可以通过等差数列求和公式$O(1)$计算。

实现代码

def solution(n: int, m: int, p: int) -> int:
    def calc(i, v):
        if v > i:
            return (v + v - i + 1) * i // 2
        return v * (v + 1) // 2 + i - v

    def check(x):
        return calc(p, x) + calc(n - p + 1, x) - x

    l, r = 1, m
    while l <= r:
        mid = l + r >> 1
        if check(mid) <= m:
            l = mid + 1
        else:
            r = mid - 1
    return r

if __name__ == '__main__':
    print(solution(4, 5, 2) == 2)
    print(solution(3, 7, 1) == 3)
    print(solution(5, 9, 3) == 3)

复杂度分析

• 时间复杂度：

  • 二分查找的时间复杂度为 O(log m)。每次 check 调用的时间复杂度为 O(1)，因为 calc 函数的计算是常数时间。因此，总的时间复杂度为 O(log m)。

• 空间复杂度：

  • 使用了常数级别的额外空间，空间复杂度为 O(1)。