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

题目描述

这是 LeetCode 上的 1802. 有界数组中指定下标处的最大值 ，难度为 中等。

给你三个正整数 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] 满足所有条件。不存在其他在指定下标处具有更大值的有效数组。

示例 2：

输入：n = 6, index = 1,  maxSum = 10

输出：3

提示：

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

---

二分 + 贪心 + 数学

根据题意，容易想到以 ans 为分割点的正整数数组具有二段性，其中 ans 为最大的 $nums[idx]$。

小于等于 ans 的值均能通过直接调整 $nums[idx]$ 来构造，不会违反总和不超过 max 的限制；大于 ans 的值则无法满足 max 限制。基于此我们可通过「二分」的方式来找分割点。

假设当前二分到的值为 x，考虑如何实现一个 check 函数，该函数用于判断 x 能否作为 $nums[idx]$：

为了令 $nums[idx] = x$ 时，数组总和 sum 不超过 max 限制，我们应当贪心构造 $nums$ 的剩余元素：从 $idx$ 开始往两侧构造，按照递减的方式进行逐个构造，若递减到 $1$ 则维持不变。

这样可确保构造出来的 $nums$ 既满足 $nums[idx] = x$ 同时元素总和最小。

位置 idx 的值为 x，其左边有 idx 个元素，其右边有 n - idx - 1 个元素。

利用「等差数列求和」公式分别从 x - 1 开始构造（注意：这里说的构造仅是计算 $nums$ 总和），若总和不超过 max 说明 $nums[idx] = x$ 满足要求，我们令 $l = mid$，否则令 $r = mid - 1$。

Java 代码：

class Solution {
    public int maxValue(int n, int index, int max) {
        long l = 1, r = max;
        while (l < r) {
            long mid = l + r + 1 >> 1;
            if (check(n, mid, index, max)) l = mid;
            else r = mid - 1;
        }
        return (int) r;
    }
    boolean check(int n, long x, int idx, int max) {
        long sum = x;
        if (idx > x - 1) {
            long an = x - 1, a1 = 1, cnt = x - 1;
            sum += cnt * (a1 + an) / 2;
            sum += idx - cnt;
        } else {
            long cnt = idx, an = x - 1, a1 = an - cnt + 1;
            sum += cnt * (a1 + an) / 2;
        }
        if (n - idx - 1 > x - 1) {
            long an = x - 1, a1 = 1, cnt = x - 1;
            sum += cnt * (a1 + an) / 2;
            sum += n - idx - 1 - cnt;
        } else {
            long cnt = n - idx - 1, an = x - 1, a1 = an - cnt + 1;
            sum += cnt * (a1 + an) / 2;
        }
        return sum <= max;
    }
}

C++ 代码：

class Solution {
public:
    int maxValue(int n, int index, int maxSum) {
        long l = 1, r = maxSum;
        while (l < r) {
            long mid = (l + r + 1) >> 1;
            if (check(n, mid, index, maxSum)) l = mid;
            else r = mid - 1;
        }
        return (int) r;
    }
    bool check(int n, long x, int idx, int maxSum) {
        long sum = x;
        if (idx > x - 1) {
            long an = x - 1, a1 = 1, cnt = x - 1;
            sum += cnt * (a1 + an) / 2;
            sum += idx - cnt;
        } else {
            long cnt = idx, an = x - 1, a1 = an - cnt + 1;
            sum += cnt * (a1 + an) / 2;
        }
        if (n - idx - 1 > x - 1) {
            long an = x - 1, a1 = 1, cnt = x - 1;
            sum += cnt * (a1 + an) / 2;
            sum += n - idx - 1 - cnt;
        } else {
            long cnt = n - idx - 1, an = x - 1, a1 = an - cnt + 1;
            sum += cnt * (a1 + an) / 2;
        }
        return sum <= maxSum;
    }
};

Python 代码：

class Solution:
    def maxValue(self, n: int, index: int, maxSum: int) -> int:
        l, r = 1, maxSum
        while l < r:
            mid = (l + r + 1) >> 1
            if self.check(n, mid, index, maxSum): 
                l = mid
            else: 
                r = mid - 1
        return int(r)

    def check(self, n, x, idx, maxSum):
        sumv = x
        if idx > x - 1:
            an, a1, cnt = x - 1, 1, x - 1
            sumv += cnt * (a1 + an) // 2
            sumv += idx - cnt
        else:
            cnt, an, a1 = idx, x - 1, x - idx
            sumv += cnt * (a1 + an) // 2
        if n - idx - 1 > x - 1:
            an, a1, cnt = x - 1, 1, x - 1
            sumv += cnt * (a1 + an) // 2
            sumv += n - idx - 1 - cnt
        else:
            cnt, an, a1 = n - idx - 1, x - 1, x - n + idx + 1
            sumv += cnt * (a1 + an) // 2
        return sumv <= maxSum

TypeScript 代码：

function maxValue(n: number, index: number, maxSum: number): number {
    const check = function(n: number, x: number, idx: number, maxSum: number): boolean {
        let sum = x;
        if (idx > x - 1) {
            let an = x - 1, a1 = 1, cnt = x - 1;
            sum += cnt * (a1 + an) / 2;
            sum += idx - cnt;
        } else {
            let cnt = idx, an = x - 1, a1 = an - cnt + 1;
            sum += cnt * (a1 + an) / 2;
        }
        if (n - idx - 1 > x - 1) {
            let an = x - 1, a1 = 1, cnt = x - 1;
            sum += cnt * (a1 + an) / 2;
            sum += n - idx - 1 - cnt;
        } else {
            let cnt = n - idx - 1, an = x - 1, a1 = an - cnt + 1;
            sum += cnt * (a1 + an) / 2;
        }
        return sum <= maxSum;
    };
    let l = 1, r = maxSum;
    while (l < r) {
        let mid = (l + r + 1) >> 1;
        if (check(n, mid, index, maxSum)) l = mid;
        else r = mid - 1;
    }
    return r;
};

• 时间复杂度：$O(\log{n})$
• 空间复杂度：$O(1)$

---

最后

这是我们「刷穿 LeetCode」系列文章的第 No.1802 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。