leetcode - 1802. Maximum Value at a Given Index in a Bounded Array

Description

You are given three positive integers: n, index, and maxSum. You want to construct an array nums (0-indexed) that satisfies the following conditions:

nums.length == n
nums[i] is a positive integer where 0 <= i < n.
abs(nums[i] - nums[i+1]) <= 1 where 0 <= i < n-1.
The sum of all the elements of nums does not exceed maxSum.
nums[index] is maximized.
Return nums[index] of the constructed array.

Note that abs(x) equals x if x >= 0, and -x otherwise.

Example 1:

Input: n = 4, index = 2,  maxSum = 6
Output: 2
Explanation: nums = [1,2,2,1] is one array that satisfies all the conditions.
There are no arrays that satisfy all the conditions and have nums[2] == 3, so 2 is the maximum nums[2].

Example 2:

Input: n = 6, index = 1,  maxSum = 10
Output: 3

Constraints:

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

Solution

Solved after help…

From the description, if we want to make nums[index] the largest, assume it’s target, then the constructed array should be: [..., target - 3, target - 2, target - 1, target, target - 1, target - 2,...]

So we could calculate the sum of the array, and use binary search to find the target.

Note that if target is small, then we need to keep all the numbers in array be positive.

Time complexity: $o(\log n)$
Space complexity: $o(1)$

Code

class Solution:
    def maxValue(self, n: int, index: int, maxSum: int) -> int:
        def get_sum(target: int) -> int:
            if target - index >= 1:
                sum_prefix = (target - index + target) * (index + 1) // 2
            else:
                sum_prefix = (1 + target) * target // 2 + index - target + 1
            if target - (n - index - 1) >= 1:
                sum_suffix = (target + target - (n - index - 1)) * (n - index) // 2
            else:
                sum_suffix = (target + 1) * target // 2 + (n - index - target)
            sum_of_list = sum_prefix + sum_suffix - target
            return sum_of_list
        left, right = 1, 1000000000
        while left < right:
            mid = (left + right + 1) >> 1
            if get_sum(mid) > maxSum:
                right = mid - 1
            else:
                left = mid
        return (left + right + 1) >> 1