LeetCode Solutions

813. Largest Sum of Averages

Time: $O(Kn^2)$

Space: $O(Kn)$

			

class Solution {
 public:
  double largestSumOfAverages(vector<int>& nums, int K) {
    const int n = nums.size();
    // dp[i][k] := largest score to partition first i nums into k groups
    dp.resize(n + 1, vector<double>(K + 1));
    prefix.resize(n + 1);

    partial_sum(begin(nums), end(nums), begin(prefix) + 1);
    return largestSumOfAverages(nums, n, K);
  }

 private:
  vector<vector<double>> dp;
  vector<double> prefix;

  double largestSumOfAverages(const vector<int>& A, int i, int k) {
    if (k == 1)
      return prefix[i] / i;
    if (dp[i][k])
      return dp[i][k];

    // Try all possible partitions
    for (int j = k - 1; j < i; ++j)
      dp[i][k] = max(dp[i][k], largestSumOfAverages(A, j, k - 1) +
                                   (prefix[i] - prefix[j]) / (i - j));

    return dp[i][k];
  }
};
			

class Solution {
  public double largestSumOfAverages(int[] A, int K) {
    final int n = A.length;
    // dp[i][k] := largest score to partition first i nums into k groups
    dp = new double[n + 1][K + 1];
    prefix = new double[n + 1];

    for (int i = 0; i < n; ++i)
      prefix[i + 1] = A[i] + prefix[i];

    return largestSumOfAverages(A, n, K);
  }

  private double[][] dp;
  private double[] prefix;

  private double largestSumOfAverages(int[] A, int i, int k) {
    if (k == 1)
      return prefix[i] / i;
    if (dp[i][k] > 0.0)
      return dp[i][k];

    // Try all possible partitions
    for (int j = k - 1; j < i; ++j)
      dp[i][k] =
          Math.max(dp[i][k], largestSumOfAverages(A, j, k - 1) + (prefix[i] - prefix[j]) / (i - j));

    return dp[i][k];
  }
}
			

class Solution {
 public:
  double largestSumOfAverages(vector<int>& nums, int K) {
    const int n = nums.size();
    // dp[i][k] := largest score to partition first i nums into k groups
    vector<vector<double>> dp(n + 1, vector<double>(K + 1));
    vector<double> prefix(n + 1);

    partial_sum(begin(nums), end(nums), begin(prefix) + 1);

    for (int i = 1; i <= n; ++i)
      dp[i][1] = prefix[i] / i;

    for (int k = 2; k <= K; ++k)
      for (int i = k; i <= n; ++i)
        for (int j = k - 1; j < i; ++j) {
          const double average = (prefix[i] - prefix[j]) / (i - j);
          dp[i][k] = max(dp[i][k], dp[j][k - 1] + average);
        }

    return dp[n][K];
  }
};