Python o(n) time, o(1) space. Greedy solution.


  • 16
    J

    A greedy solution using three sliding windows where you keep track of the best indexes/sums as you go.

    O(n) time: Since we're only going through the list once and using no complex operations, this is O(n).
    O(1) space: Just a fixed set of temp vars. We don't need the extra arrays that the DP solutions have.

    class Solution:
        def maxSumOfThreeSubarrays(self, nums, k):
            """
            :type nums: List[int]
            :type k: int
            :rtype: List[int]
            """
    
            # Best single, double, and triple sequence found so far
            bestSeq = 0
            bestTwoSeq = [0, k]
            bestThreeSeq = [0, k, k*2]
    
            # Sums of each window
            seqSum = sum(nums[0:k])
            seqTwoSum = sum(nums[k:k*2])
            seqThreeSum = sum(nums[k*2:k*3])
    
            # Sums of combined best windows
            bestSeqSum = seqSum
            bestTwoSum = seqSum + seqTwoSum
            bestThreeSum = seqSum + seqTwoSum + seqThreeSum
    
            # Current window positions
            seqIndex = 1
            twoSeqIndex = k + 1
            threeSeqIndex = k*2 + 1
            while threeSeqIndex <= len(nums) - k:
                # Update the three sliding windows
                seqSum = seqSum - nums[seqIndex - 1] + nums[seqIndex + k - 1]
                seqTwoSum = seqTwoSum - nums[twoSeqIndex - 1] + nums[twoSeqIndex + k - 1]
                seqThreeSum = seqThreeSum - nums[threeSeqIndex - 1] + nums[threeSeqIndex + k - 1]
                
                # Update best single window
                if seqSum > bestSeqSum:
                    bestSeq = seqIndex
                    bestSeqSum = seqSum
    
                # Update best two windows
                if seqTwoSum + bestSeqSum > bestTwoSum:
                    bestTwoSeq = [bestSeq, twoSeqIndex]
                    bestTwoSum = seqTwoSum + bestSeqSum
    
                # Update best three windows
                if seqThreeSum + bestTwoSum > bestThreeSum:
                    bestThreeSeq = bestTwoSeq + [threeSeqIndex]
                    bestThreeSum = seqThreeSum + bestTwoSum
    
                # Update the current positions
                seqIndex += 1
                twoSeqIndex += 1
                threeSeqIndex += 1
    
            return bestThreeSeq
    

  • 1
    O

    Many people solved this problem using DP. But just three pointer is enough. It is quite good and underrated solution :)


  • 2

    @JaRail
    Thanks for your brilliant solution. Here is a C++ version solution according to your idea.

        vector<int> maxSumOfThreeSubarrays(vector<int>& nums, int k) {
            int n = nums.size();
            // sum[k]: records the current kth subarray sum
            // max_sum[k]: records the max sum of 0th, ..., kth subarray sum 
            vector<int> sum(3, 0), max_sum(3, 0);
            vector<vector<int>> idx(3, vector<int>(3, 0));
            
            // records the index
            idx[0] = vector<int>{0};
            idx[1] = vector<int>{0, k};
            idx[2] = vector<int>{0, k, 2*k};
            
            // init the sum
            sum[0] = accumulate(nums.begin(), nums.begin() + k, 0);
            sum[1] = accumulate(nums.begin() + k, nums.begin() + 2*k, 0);
            sum[2] = accumulate(nums.begin() + 2*k, nums.begin() + 3*k, 0);
            
            // init the max_sum
            max_sum[0] = sum[0];
            max_sum[1] = sum[0] + sum[1];
            max_sum[2] = sum[0] + sum[1] + sum[2];
            
            for (int i = 1; i <= n - 3*k; ++i) {
                // update current sum
                sum[0] += nums[i+k-1] - nums[i-1];
                sum[1] += nums[i+2*k-1] - nums[i+k-1];
                sum[2] += nums[i+3*k-1] - nums[i+2*k-1];
                
                // update max_sum[0] and idx[0] if possible
                if (sum[0] > max_sum[0]) {
                    max_sum[0] = sum[0];
                    idx[0][0] = i;
                }
                
                // update max_sum[1] and idx[1] if possible
                if (sum[1] + max_sum[0] > max_sum[1]) {
                    max_sum[1] = sum[1] + max_sum[0];
                    idx[1][0] = idx[0][0];
                    idx[1][1] = i+k;
                }
                
                // update max_sum[2] and idx[2] if possible
                if (sum[2] + max_sum[1] > max_sum[2]) {
                    max_sum[2] = sum[2] + max_sum[1];
                    idx[2][0] = idx[1][0];
                    idx[2][1] = idx[1][1];
                    idx[2][2] = i + 2*k;
                }
            }
            
            return idx[2];
        }
    

  • 0
    J

    @yanchao_hust Nice. It looks a lot cleaner with the vectors. I was thinking of doing a generic solution for m subarrays. That would clean it up a bit more. But mostly I just wanted to show an o(1) space solution was reasonable. Thanks for your contribution!


  • 0

    @JaRail
    Thanks for your reply. I will give it a try for m subarray general case.


  • 1

    @JaRail
    Now for any m subarray.

    vector<int> maxSumOfThreeSubarrays(vector<int>& nums, int k) {
            int n = nums.size();
            int m = 3;
            // sum[k]: records the current kth subarray sum
            // max_sum[k]: records the max sum of 0th, ..., kth subarray sum 
            vector<int> sum(m + 1, 0), max_sum(m + 1, 0);
            vector<vector<int>> idx(m);
            
            // records the index
            for (int i = 0; i < m; ++i) {
                for (int j = 0; j <= i; ++j) {
                    idx[i].push_back(j*k);
                }
                
                // init the sum
                sum[i + 1] = accumulate(nums.begin() + k*i, nums.begin() + (i+1)*k, 0);
                max_sum[i + 1] = max_sum[i] + sum[i + 1];
                
            }
            
            
            for (int i = 1; i <= n - m*k; ++i) {
                // update current sum
                for (int j = 0; j < m; ++j) {
                    sum[j + 1] += nums[i + (j+1)*k - 1] - nums[i + j*k - 1];
                    if (sum[j + 1] + max_sum[j] > max_sum[j + 1]) {
                        max_sum[j + 1] = sum[j + 1] + max_sum[j];
                        // update index
                        idx[j][j] = i + j*k;
                        for (int k = 0; k < j; ++k) {
                            idx[j][k] = idx[j-1][k];
                        } 
                    }
                }
            }
            
            return idx[m-1];
        }
    

  • 0
    J

    @yanchao_hust beautiful! :)


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.