3427-SumOfVariableLengthSubarrays.go

package main

// 3427. Sum of Variable Length Subarrays
// You are given an integer array nums of size n. 
// For each index i where 0 <= i < n, define a subarray nums[start ... i] where start = max(0, i - nums[i]).

// Return the total sum of all elements from the subarray defined for each index in the array.

// Example 1:
// Input: nums = [2,3,1]
// Output: 11
// Explanation:
// i	Subarray	Sum
// 0	nums[0] = [2]	2
// 1	nums[0 ... 1] = [2, 3]	5
// 2	nums[1 ... 2] = [3, 1]	4
// Total Sum	 	11
// The total sum is 11. Hence, 11 is the output.

// Example 2:
// Input: nums = [3,1,1,2]
// Output: 13
// Explanation:
// i	Subarray	Sum
// 0	nums[0] = [3]	3
// 1	nums[0 ... 1] = [3, 1]	4
// 2	nums[1 ... 2] = [1, 1]	2
// 3	nums[1 ... 3] = [1, 1, 2]	4
// Total Sum	 	13
// The total sum is 13. Hence, 13 is the output.

// Constraints:
//     1 <= n == nums.length <= 100
//     1 <= nums[i] <= 1000

import "fmt"

func subarraySum(nums []int) int {
    res, sum, n := 0, 0, len(nums)
    prefix := make([]int, n)
    for i, v := range nums {
        sum += v
        prefix[i] = sum
    }
    max := func (x, y int) int { if x > y { return x; }; return y; }
    for i, v := range nums {
        start := max(0, i - v)
        res += prefix[i] - (prefix[start] - nums[start])
    }
    return res
}

func subarraySum1(nums []int) int {
    res, n := 0, len(nums)
    prefix := make([]int, n + 1)
    for i, v := range nums {
        prefix[i + 1] = prefix[i] + v
    }
    max := func (x, y int) int { if x > y { return x; }; return y; }
    for i, v := range nums {
        res += prefix[i + 1] - prefix[max(0, i - v)]
    }
    return res
}

func main() {
    // Example 1:
    // Input: nums = [2,3,1]
    // Output: 11
    // Explanation:
    // i	Subarray	Sum
    // 0	nums[0] = [2]	2
    // 1	nums[0 ... 1] = [2, 3]	5
    // 2	nums[1 ... 2] = [3, 1]	4
    // Total Sum	 	11
    // The total sum is 11. Hence, 11 is the output.
    fmt.Println(subarraySum([]int{2,3,1})) // 11
    // Example 2:
    // Input: nums = [3,1,1,2]
    // Output: 13
    // Explanation:
    // i	Subarray	Sum
    // 0	nums[0] = [3]	3
    // 1	nums[0 ... 1] = [3, 1]	4
    // 2	nums[1 ... 2] = [1, 1]	2
    // 3	nums[1 ... 3] = [1, 1, 2]	4
    // Total Sum	 	13
    // The total sum is 13. Hence, 13 is the output.
    fmt.Println(subarraySum([]int{3,1,1,2})) // 13

    fmt.Println(subarraySum([]int{1,2,3,4,5,6,7,8,9})) // 165
    fmt.Println(subarraySum([]int{9,8,7,6,5,4,3,2,1})) // 175

    fmt.Println(subarraySum1([]int{2,3,1})) // 11
    fmt.Println(subarraySum1([]int{3,1,1,2})) // 13
    fmt.Println(subarraySum1([]int{1,2,3,4,5,6,7,8,9})) // 165
    fmt.Println(subarraySum1([]int{9,8,7,6,5,4,3,2,1})) // 175
}