LC-1802.有界数组中指定下标处的最大值

题目描述

leetcode 中等题

给你三个正整数 n、index 和 maxSum 。你需要构造一个同时满足下述所有条件的数组 nums（下标 从 0 开始 计数）：

• nums.length == n
• nums[i] 是 正整数 ，其中 0 <= i < n
• abs(nums[i] - nums[i+1]) <= 1 ，其中 0 <= i < n-1
• nums 中所有元素之和不超过 maxSum
• nums[index] 的值被 最大化

返回你所构造的数组中的 nums[index] 。

注意：abs(x) 等于 x 的前提是 x >= 0 ；否则，abs(x) 等于 -x 。

示例1：

输入：n = 4, index = 2,  maxSum = 6
输出：2
解释：数组 [1,1,2,1] 和 [1,2,2,1] 满足所有条件。不存在其他在指定下标处具有更大值的有效数组。

提示1：

1 <= n <= maxSum <= 10^9
0 <= index < n

二分查找

首先根据题目描述，为了使 $index$ 处元素值尽量大，显然 $index$ 处的两侧元素必然是从 $nums[index]-1$ 递减到 $1$ 的，并且如果某一侧递减到 $1$ 后还存在剩余位置，就用 $1$ 补齐。

由此可以进行数组的元素总和的计算，设 $index$ 处元素值为 $mid$，某一侧的元素数量为 $cnt$

• 当 $cnt > mid$ 时，该侧的元素和为 $\frac{mid \times (1+mid) }{2} + (cnt - mid)$，其中前者为 $1$ 到 $mid$ 的等差数列和，后者为需要填充 $1$ 的数量。
• 当 $cnt <= mid$ 时，该侧的元素和就为 $\frac{cnt \times (mid - cnt + 1 +mid) }{2}$，即从 $mid - cnt + 1$ 到 $mid$ 的等差数列和。

由于题目 $maxSum$ 范围为 $10^9$，所以 $O(maxSum)$ 的遍历会 TLE。又因为随着 $mid$ 的增大，数组的总和也会增大，所以可以使用二分查找右边界的方式来找到最后一个满足小于等于 $maxSum$ 的值，此时的时间复杂度为 $O(log(maxSum))$。

class Solution {
    public int maxValue(int n, int index, int maxSum) {
        int left = 1, right = maxSum + 1;
        while(left < right){
            int mid = left + (right - left >> 1);
            if(getSum(mid, index + 1) + getSum(mid - 1, n - index - 1) <= maxSum){
                left = mid + 1;
            }else{
                right = mid;
            }
        }
        return left - 1;
    }

    private long getSum(long mid, long cnt){
        if(cnt > mid){
            return (mid * (1 + mid)) / 2 + (cnt - mid);
        }
        return (cnt * (mid - cnt + 1 + mid)) / 2;
    }
}