Stone Game II Interview Question for Amazon

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

// Naive recursive solution (for understanding, not efficient enough for larger inputs)
int stoneGameII_recursive(const vector<int>& piles, int i, int M, vector<int>& prefixSum)
{
    if (i >= piles.size()) {
        return 0;
    }

    int remaining = piles.size() - i;
    int maxStones = 0;

    for (int X = 1; X <= 2 * M && X <= remaining; ++X) {
        int currentStones = prefixSum[i + X] - prefixSum[i];
        
        // The other player's score is total - our score, we want to maximize our score so minimize the score of the other player
        int nextM = max(M, X);
        int opponentStones = stoneGameII_recursive(piles, i + X, nextM, prefixSum);
        maxStones = max(maxStones, currentStones - opponentStones);
    }
    
    return maxStones;
}

// Optimized dynamic programming solution
int stoneGameII_dp(const vector<int>& piles)
{
    int n = piles.size();
    vector<int> prefixSum(n + 1, 0);
    for (int i = 0; i < n; ++i) {
        prefixSum[i + 1] = prefixSum[i] + piles[i];
    }

    // dp[i][M] represents the maximum difference between Alice and Bob's score
    // starting from index i with current M
    vector<vector<int>> dp(n, vector<int>(n, 0));

    // Iterate backwards to fill dp table (bottom-up)
    for (int i = n - 1; i >= 0; --i) {
        for (int M = n; M >= 1; --M) {
            if (i + 2 * M >= n) {
                // Can take all remaining stones
                dp[i][M] = prefixSum[n] - prefixSum[i];
            } else {
                // Iterate through possible values of X
                for (int X = 1; X <= 2 * M; ++X) {
                    if (i + X <= n) {
                        int nextM = max(M, X);
                        dp[i][M] = max(dp[i][M], (prefixSum[i + X] - prefixSum[i]) - dp[i + X][nextM]);
                    }
                }
            }
        }
    }

    // The total stones are sum of the array, 
    // so Alice score = (total stones + difference) / 2
    return (prefixSum[n] + dp[0][1]) / 2;
}

int main()
{
    vector<int> piles1 = {2, 7, 9, 4, 4};
    cout << "Maximum stones Alice can get (DP): " << stoneGameII_dp(piles1) << endl;
    
    vector<int> piles2 = {1, 2, 3, 4, 5, 100};
    cout << "Maximum stones Alice can get (DP): " << stoneGameII_dp(piles2) << endl;
    
    vector<int> piles3 = {2, 7, 9, 4, 4};
    vector<int> prefixSum3(piles3.size() + 1, 0);
    for (int i = 0; i < piles3.size(); ++i) {
        prefixSum3[i + 1] = prefixSum3[i] + piles3[i];
    }
    cout << "Maximum stones Alice can get (Recursive): " << stoneGameII_recursive(piles3, 0, 1, prefixSum3) << endl;

    return 0;
}

Explanation:

Naive Recursive Solution:

Base Case: If i is out of bounds (i.e., i >= piles.size()), return 0 because there are no more stones to take.
Recursive Step:
- Iterate through all possible values of X (from 1 to 2*M), ensuring X does not exceed the remaining piles.
- Calculate the current stones taken by Alice: currentStones = prefixSum[i + X] - prefixSum[i]. The prefix sum helps to get sum of stones in O(1).
- Recursively calculate the opponent's best score by calling stoneGameII_recursive for the remaining piles.
- Update maxStones to maximize Alice's score: maxStones = max(maxStones, currentStones - opponentStones).

Optimized Dynamic Programming Solution:

Initialization:
- Calculate prefixSum array to efficiently compute sum of stones.
- Initialize dp[n][n] table with zeros. dp[i][M] represents the maximum difference Alice can achieve starting from index i with current M.
DP Iteration:
- Iterate backwards starting from the last pile to the first (i = n - 1 to 0).
- For each pile, iterate through possible values of M (from n down to 1).
- If i + 2 * M >= n, Alice can take all remaining stones, so dp[i][M] = prefixSum[n] - prefixSum[i].
- Otherwise, iterate through all possible X (from 1 to 2 * M):
  - Update dp[i][M] with the maximum difference Alice can get: dp[i][M] = max(dp[i][M], (prefixSum[i + X] - prefixSum[i]) - dp[i + X][nextM]), where nextM = max(M, X).
Final Result:
- The maximum stones Alice can get is calculated as (prefixSum[n] + dp[0][1]) / 2.

Big(O) Run-time Analysis:

Naive Recursive Solution:
- Time Complexity: O(2^N) in the worst case, because for each state, there can be multiple branches to explore. However, it's an approximation. It's more like O((2M)^N), since there are at most 2M options at each step.
- Reasoning: The function branches into multiple recursive calls, leading to exponential time complexity. Each call explores possibilities of taking 1 to 2M piles.
Optimized Dynamic Programming Solution:
- Time Complexity: O(N^2 * M), where N is the number of piles and M is capped by N. Since M can grow up to N, effectively it is O(N^3).
- Reasoning: The algorithm uses nested loops to fill the DP table. The outer loops iterate from n-1 to 0 and from n to 1. The inner loop iterates up to 2*M which is at most 2*N. Thus the overall complexity is O(N * N * N).

Big(O) Space Usage Analysis:

Naive Recursive Solution:
- Space Complexity: O(N) in the worst case due to the call stack. This is because the depth of the recursion can go up to N in the worst scenario (e.g., if the values of M always allow taking only one pile at a time).
- Reasoning: The call stack grows proportionally to the depth of the recursion, which is influenced by the number of piles.
Optimized Dynamic Programming Solution:
- Space Complexity: O(N^2) because of the dp[n][n] table. Also, O(N) for the prefix sum array. Thus the overall space is O(N^2 + N) -> O(N^2)
- Reasoning: The DP table stores the intermediate results for each possible state, requiring a 2D array of size n x n.

Edge Cases:

Empty Input:
- If piles is empty, return 0. Add a check at the beginning of the function to handle this case.
Single Pile:
- If there is only one pile, Alice will take it, and the result should be the value of that pile.
All Piles are Zero:
- If all piles[i] are zero, the maximum stones Alice can get is zero.
Large Input:
- For large inputs, the dynamic programming solution is more efficient than the recursive solution.
Constraints on M:
- The maximum value of M is constrained by the number of remaining piles.