Leetcode-1802

Maximum Value at a Given Index in a Bounded Array

Time complexity: $O(1)$

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
#include <iostream>
#include <cmath>

class Solution {
public:
    int maxValue(int n, int idx, int maxSum) {
        idx = std::max(idx, n - idx - 1);
        long long R = 2 * idx + 1 - n;
        long long h = idx + 1;
        long long sum = h * h - (R * (R + 1)) / 2;
        if (maxSum >= sum) {
            return h + (maxSum - sum) / n;
        }
        
        h = n - idx;
        sum = h * h + (idx - h + 1);
        if (sum > maxSum) {
            return 1 + (int)std::sqrt(maxSum - n);
        }
        else {
            // r=h+idx-n
            // h^2+(idx-h+1)-r(r+1)/2<=maxSum
            long long b = 2 * n - 2 * idx - 3;
            long long c = 2ll * idx + 2 - (long long)(idx - n) * (idx - n + 1) - 2ll * maxSum;
            return (int)((-b + (long long)std::sqrt(b * b - 4 * c)) / 2);
        }
    }
};

Leetcode-1802

Licensed under CC BY-NC-SA 4.0