2607-MakeKSubarraySumsEqual.go

package main

// 2607. Make K-Subarray Sums Equal
// You are given a 0-indexed integer array arr and an integer k. 
// The array arr is circular. 
// In other words, the first element of the array is the next element of the last element, 
// and the last element of the array is the previous element of the first element.

// You can do the following operation any number of times:
//     Pick any element from arr and increase or decrease it by 1.

// Return the minimum number of operations such that the sum of each subarray of length k is equal.

// A subarray is a contiguous part of the array.

// Example 1:
// Input: arr = [1,4,1,3], k = 2
// Output: 1
// Explanation: we can do one operation on index 1 to make its value equal to 3.
// The array after the operation is [1,3,1,3]
// - Subarray starts at index 0 is [1, 3], and its sum is 4 
// - Subarray starts at index 1 is [3, 1], and its sum is 4 
// - Subarray starts at index 2 is [1, 3], and its sum is 4 
// - Subarray starts at index 3 is [3, 1], and its sum is 4 

// Example 2:
// Input: arr = [2,5,5,7], k = 3
// Output: 5
// Explanation: we can do three operations on index 0 to make its value equal to 5 and two operations on index 3 to make its value equal to 5.
// The array after the operations is [5,5,5,5]
// - Subarray starts at index 0 is [5, 5, 5], and its sum is 15
// - Subarray starts at index 1 is [5, 5, 5], and its sum is 15
// - Subarray starts at index 2 is [5, 5, 5], and its sum is 15
// - Subarray starts at index 3 is [5, 5, 5], and its sum is 15 

// Constraints:
//     1 <= k <= arr.length <= 10^5
//     1 <= arr[i] <= 10^9

import "fmt"
import "sort"

func makeSubKSumEqual(arr []int, k int) int64 {
    abs := func(x int) int { if x < 0 { return -x; }; return x; }
    gcd := func (x, y int) int { for y != 0 { x, y = y, x % y; }; return x; }
    res, n := 0, len(arr)
    steppedListNum := gcd(n, k)
    steppedListLen := n / steppedListNum
    steppedList := make([]int, steppedListLen)
    for shift := 0; shift < steppedListNum; shift++ {
        for i := 0; i < steppedListLen; i++ {
            steppedList[i] = arr[i * steppedListNum + shift]
        }
        sort.Ints(steppedList)
        median := steppedList[steppedListLen / 2]
        for _, v := range steppedList {
            res += abs(median - v)
        }
    }
    return int64(res)
}

func makeSubKSumEqual1(arr []int, k int) int64 {
    abs := func(x int) int { if x < 0 { return -x; }; return x; }
    gcd := func (x, y int) int { for y != 0 { x, y = y, x % y; }; return x; }
    k = gcd(k, len(arr))
    res, g := 0, make([][]int, k)
    for i, v := range arr {
        g[i % k] = append(g[i % k], v)
    }
    for _, row := range g {
        sort.Ints(row)
        for _, v := range row {
            res += abs(v - row[len(row) / 2])
        }
    }
    return int64(res)
}

func main() {
    // Example 1:
    // Input: arr = [1,4,1,3], k = 2
    // Output: 1
    // Explanation: we can do one operation on index 1 to make its value equal to 3.
    // The array after the operation is [1,3,1,3]
    // - Subarray starts at index 0 is [1, 3], and its sum is 4 
    // - Subarray starts at index 1 is [3, 1], and its sum is 4 
    // - Subarray starts at index 2 is [1, 3], and its sum is 4 
    // - Subarray starts at index 3 is [3, 1], and its sum is 4 
    fmt.Println(makeSubKSumEqual([]int{1,4,1,3}, 2)) // 1
    // Example 2:
    // Input: arr = [2,5,5,7], k = 3
    // Output: 5
    // Explanation: we can do three operations on index 0 to make its value equal to 5 and two operations on index 3 to make its value equal to 5.
    // The array after the operations is [5,5,5,5]
    // - Subarray starts at index 0 is [5, 5, 5], and its sum is 15
    // - Subarray starts at index 1 is [5, 5, 5], and its sum is 15
    // - Subarray starts at index 2 is [5, 5, 5], and its sum is 15
    // - Subarray starts at index 3 is [5, 5, 5], and its sum is 15 
    fmt.Println(makeSubKSumEqual([]int{2,5,5,7}, 2)) // 5

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

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