Start learning free

LESSON 231 OF 410 · PARTIAL SUMS #1

Maximum Sum Subarray of Size K in Python: Brute Force, O(n * k)

Practice O(n * k) Read the full explanation

Given an input array of integers and a positive integer k, find the contiguous subarray of length exactly k with the maximum sum, and return its first and last indices.

Example

Input: nums = [5, 6, 1, 2, 6, 6, 4, 3], k = 3
Output: [4, 6]
Explanation: The subarray nums[4...6] has the maximum sum of 16
STEP 1 OF 8
    Ctrl + Enter to check

    PYTHON · PARTIAL SUMS #1

    Maximum Sum Subarray Of Length K in Python, explained

    Problem overview

    The method maxSum receives a list of integers, nums, and a positive integer k. A subarray is a run of neighbouring elements, and this one must have length exactly k. Among all such windows, the method finds the one with the largest sum and returns its first and last index as a list [start, end], not the sum itself.

    Take nums = [5, 6, 1, 2, 6, 6, 4, 3] and k = 3. There are six windows of length 3:

    StartElementsSum
    05, 6, 112
    16, 1, 29
    21, 2, 69
    32, 6, 614
    46, 6, 416
    56, 4, 313

    The best sum is 16, from the window that starts at index 4 and ends at index 6, so the method returns [4, 6]. If two windows tie, this lesson keeps the earlier one.

    Approach: try every window

    The most direct way to find the maximum sum subarray of size k is to check every window. A list of n elements has n - k + 1 windows of length k, because the last one has to start early enough to fit: its start index is len(nums) - k. For each window, add up its k elements, and remember the best sum together with the window's bounds.

    That is k additions for each of n - k + 1 windows, so the time is O(n * k). It is simple and correct, and it is a good first version. The faster idea, the sliding window, notices that neighbouring windows share k - 1 elements: moving the window one step to the right adds one number and removes one, so each sum comes from the previous sum in constant time. That gives O(n). This lesson builds the brute force version, and the follow-up lesson on the same problem improves it.

    Step by step

    1. Start the best sum at negative infinity. maxSum = float('-inf') is smaller than any real sum, so the first window always replaces it. Starting at 0 would be a bug when every window has a negative sum.
    2. Create an empty list answer for the best window's bounds.
    3. Loop over window starts. range(len(nums) - k + 1) gives 0 up to len(nums) - k. The end value of range is excluded, which is why the + 1 is there.
    4. Reset the sum for each window. sum = 0 sits inside the outer loop, so every start adds up its own elements. Putting it before the loop would carry the previous window's total into the next one.
    5. Loop over the window's positions. range(i, i + k) visits exactly k indices, from i to i + k - 1.
    6. Add each element. sum += nums[idx].
    7. Compare after the inner loop. When the window's sum is strictly greater than maxSum, store it and set answer to [i, i + k - 1]. The strict > means a later window with an equal sum does not replace the earlier one.
    8. Return answer after both loops have finished.

    The solution in Python

    class Solution:
        def maxSum(self, nums: List[int], k: int) -> int:
            maxSum = float('-inf')
            answer = []
            for i in range(len(nums) - k + 1):
                sum = 0
                for idx in range(i, i + k):
                    sum += nums[idx]
                if sum > maxSum:
                    maxSum = sum
                    answer = [i, i + k - 1]
            return answer
    

    A few details are worth knowing:

    • The method returns a list, even though the annotation says int. Python does not check annotations when the code runs, so this works, but the return value is the pair of indices.
    • The variable sum hides Python's built-in sum() inside the method. That is harmless here because the method never calls the built-in, but in your own code a name like window_sum avoids the trap.
    • float('-inf') is a float, while the sums are ints. Python compares them without trouble.
    • Python integers have no fixed size, so the sums cannot overflow.

    Complexity

    The outer loop runs n - k + 1 times and the inner loop runs k times, so the time is O((n - k + 1) * k), usually written O(n * k). It is largest when k is around n / 2, and for k = 1 or k = n it is only O(n). The extra space is O(1): a few variables and a two-element list.

    Edge cases and common mistakes

    • All numbers negative. For [-3, -1, -2, -5] with k = 2 the sums are -4, -3 and -7, so the answer is [1, 2]. Initializing maxSum to 0 would never accept any window.
    • k equals the list length. There is exactly one window, and range(len(nums) - k + 1) is range(1). For [1, 2, 3] and k = 3 the answer is [0, 2].
    • k = 1. Every window is one element, so the result is the position of the largest number, as [i, i].
    • Ties. For [4, 4, 4] with k = 2 both windows sum to 8, and the strict > returns [0, 1]. Changing it to >= would return [1, 2].
    • k larger than the list. No window fits, the loop does not run, and the method returns the empty list [].
    • Off-by-one in the ranges. Using range(len(nums) - k) skips the last window, and range(i, i + k - 1) adds only k - 1 elements.
    • Returning the sum. The method returns the indices. The sum is only used to compare windows.

    Testing

    • nums = [5, 6, 1, 2, 6, 6, 4, 3], k = 3 returns [4, 6].
    • nums = [-3, -1, -2, -5], k = 2 returns [1, 2].
    • nums = [7, 1, 2], k = 1 returns [0, 0].
    • nums = [4, 4, 4], k = 2 returns [0, 1].
    Keep going with PythonAbout 340 step-by-step lessons, from your first line of code to coding interviews. Free to start.
    Start learning free