Split Array Largest Sum Python Solution

lcpy gen -n 410   # by problem number
lcpy gen -s split_array_largest_sum   # by problem name

Input: nums = [7,2,5,10,8], k = 2
Output: 18
Explanation: There are four ways to split nums into two subarrays.
The best way is to split it into [7,2,5] and [10,8], where the largest sum among the two subarrays is only 18.

Input: nums = [1,2,3,4,5], k = 2
Output: 9
Explanation: There are four ways to split nums into two subarrays.
The best way is to split it into [1,2,3] and [4,5], where the largest sum among the two subarrays is only 9.

class Solution:
    # Time: O(n * log(sum(nums))) binary search on the answer
    # Space: O(1)
    def split_array(self, nums: list[int], k: int) -> int:
        # Lower bound: a single element must fit. Upper bound: sum of all elements.
        low = max(nums)
        high = sum(nums)

        def count_subarrays(capacity: int) -> int:
            """Minimum subarrays needed so no subarray sum exceeds capacity."""
            subarrays = 1
            current = 0
            for value in nums:
                if current + value > capacity:
                    subarrays += 1
                    current = value
                else:
                    current += value
            return subarrays

        # Binary search for smallest capacity that fits within k subarrays.
        while low < high:
            mid = (low + high) // 2
            if count_subarrays(mid) <= k:
                high = mid
            else:
                low = mid + 1
        return low