Sum of Total Strength of Wizards Interview Question for Meta

#include <iostream>
#include <vector>

using namespace std;

const int MOD = 1e9 + 7;

// Naive Solution: O(n^2)
long long totalStrengthNaive(vector<int>& strength) {
    int n = strength.size();
    long long total_strength = 0;

    for (int i = 0; i < n; ++i) {
        for (int j = i; j < n; ++j) {
            int min_strength = strength[i];
            long long sum = 0;
            for (int k = i; k <= j; ++k) {
                min_strength = min(min_strength, strength[k]);
                sum = (sum + strength[k]) % MOD;
            }
            total_strength = (total_strength + (long long)min_strength * sum) % MOD;
        }
    }

    return total_strength;
}

// Optimized Solution: O(n)
long long totalStrengthOptimized(vector<int>& strength) {
    int n = strength.size();
    vector<long long> left(n), right(n);

    // Find left boundary: the distance to the first element smaller than strength[i] on the left
    vector<int> stack;
    for (int i = 0; i < n; ++i) {
        while (!stack.empty() && strength[stack.back()] > strength[i]) {
            stack.pop_back();
        }
        left[i] = stack.empty() ? i + 1 : i - stack.back();
        stack.push_back(i);
    }

    stack.clear();

    // Find right boundary: the distance to the first element smaller or equal to strength[i] on the right
    for (int i = n - 1; i >= 0; --i) {
        while (!stack.empty() && strength[stack.back()] >= strength[i]) {
            stack.pop_back();
        }
        right[i] = stack.empty() ? n - i : stack.back() - i;
        stack.push_back(i);
    }

    // Calculate prefix sum and prefix sum of prefix sum
    vector<long long> prefix_sum(n + 1, 0), prefix_sum_prefix_sum(n + 1, 0);
    for (int i = 0; i < n; ++i) {
        prefix_sum[i + 1] = (prefix_sum[i] + strength[i]) % MOD;
        prefix_sum_prefix_sum[i + 1] = (prefix_sum_prefix_sum[i] + prefix_sum[i + 1]) % MOD;
    }

    long long total_strength = 0;
    for (int i = 0; i < n; ++i) {
        long long left_sum = (prefix_sum_prefix_sum[i + 1] - prefix_sum_prefix_sum[max(0, i - (int)left[i] + 1)] - (long long)(i - left[i] + 1) * (prefix_sum[max(0, i - (int)left[i] + 1)])) % MOD;
        long long right_sum = ((long long)(i + right[i] + 1) * (prefix_sum[min(n, i + (int)right[i] + 1)]) - prefix_sum_prefix_sum[min(n, i + (int)right[i] + 1)] - (prefix_sum_prefix_sum[i + 1])) % MOD;
        total_strength = (total_strength + strength[i] * (left_sum * right[i] % MOD + right_sum * left[i] % MOD) % MOD) % MOD;
    }

    return (total_strength + MOD) % MOD; // Ensure result is non-negative
}

int main() {
    vector<int> strength1 = {1, 3, 1, 2};
    cout << "Total strength (naive) for {1, 3, 1, 2}: " << totalStrengthNaive(strength1) << endl;
    cout << "Total strength (optimized) for {1, 3, 1, 2}: " << totalStrengthOptimized(strength1) << endl;

    vector<int> strength2 = {5, 4, 6};
    cout << "Total strength (naive) for {5, 4, 6}: " << totalStrengthNaive(strength2) << endl;
    cout << "Total strength (optimized) for {5, 4, 6}: " << totalStrengthOptimized(strength2) << endl;

    return 0;
}

Naive Solution

The naive solution iterates through all possible subarrays, calculates the minimum strength within each subarray, calculates the sum of strengths within each subarray, and then multiplies these two values together. Finally, it sums up all these products modulo 10^9 + 7. This approach is straightforward but inefficient.

Optimal Solution

The optimized solution utilizes a stack to find the left and right boundaries for each element, indicating the range where that element is the minimum. It then uses prefix sums and prefix sums of prefix sums to efficiently calculate the sum of strengths within these ranges. This significantly reduces the time complexity.

// Optimized Solution: O(n)
long long totalStrengthOptimized(vector<int>& strength) {
    int n = strength.size();
    vector<long long> left(n), right(n);

    // Find left boundary: the distance to the first element smaller than strength[i] on the left
    vector<int> stack;
    for (int i = 0; i < n; ++i) {
        while (!stack.empty() && strength[stack.back()] > strength[i]) {
            stack.pop_back();
        }
        left[i] = stack.empty() ? i + 1 : i - stack.back();
        stack.push_back(i);
    }

    stack.clear();

    // Find right boundary: the distance to the first element smaller or equal to strength[i] on the right
    for (int i = n - 1; i >= 0; --i) {
        while (!stack.empty() && strength[stack.back()] >= strength[i]) {
            stack.pop_back();
        }
        right[i] = stack.empty() ? n - i : stack.back() - i;
        stack.push_back(i);
    }

    // Calculate prefix sum and prefix sum of prefix sum
    vector<long long> prefix_sum(n + 1, 0), prefix_sum_prefix_sum(n + 1, 0);
    for (int i = 0; i < n; ++i) {
        prefix_sum[i + 1] = (prefix_sum[i] + strength[i]) % MOD;
        prefix_sum_prefix_sum[i + 1] = (prefix_sum_prefix_sum[i] + prefix_sum[i + 1]) % MOD;
    }

    long long total_strength = 0;
    for (int i = 0; i < n; ++i) {
        long long left_sum = (prefix_sum_prefix_sum[i + 1] - prefix_sum_prefix_sum[max(0, i - (int)left[i] + 1)] - (long long)(i - left[i] + 1) * (prefix_sum[max(0, i - (int)left[i] + 1)])) % MOD;
        long long right_sum = ((long long)(i + right[i] + 1) * (prefix_sum[min(n, i + (int)right[i] + 1)]) - prefix_sum_prefix_sum[min(n, i + (int)right[i] + 1)] - (prefix_sum_prefix_sum[i + 1])) % MOD;
        total_strength = (total_strength + strength[i] * (left_sum * right[i] % MOD + right_sum * left[i] % MOD) % MOD) % MOD;
    }

    return (total_strength + MOD) % MOD; // Ensure result is non-negative
}

Big(O) Run-time Analysis

Naive Solution: The naive solution involves three nested loops, resulting in a time complexity of O(n^3). However, if we precompute sums in O(n^2) then we can optimize it to O(n^2).
Optimized Solution: The optimized solution calculates left and right boundaries in O(n) time using a stack. The prefix sum and prefix sum of prefix sum calculations also take O(n) time. The final loop iterates through the array once, performing constant-time operations, resulting in an overall time complexity of O(n).

Big(O) Space Usage Analysis

Naive Solution: The naive solution uses a constant amount of extra space, so the space complexity is O(1).
Optimized Solution: The optimized solution uses three additional vectors of size n (left, right, prefix_sum, prefix_sum_prefix_sum) resulting in a space complexity of O(n).

Edge Cases

Empty Input Array: If the input array is empty, the total strength should be 0. The provided code handles this implicitly, as the loops will not execute, and 0 will be returned.
Single Element Array: If the array contains only one element, the total strength is the square of that element. Both the naive and optimized solutions handle this case correctly.
Array with Duplicate Strengths: The algorithm correctly handles arrays with duplicate strengths, as it considers each contiguous group of wizards.
Large Strengths: Given the constraint that 1 <= strength[i] <= 10^9, it's crucial to use long long to prevent integer overflow when calculating sums and products. The provided code uses long long for these calculations.
Modulo Operation: Since the result may be very large, it's essential to apply the modulo operation (10^9 + 7) after each multiplication and addition to prevent overflow and ensure the result remains within the specified range. The provided code correctly uses the modulo operator.