package main

import (
    "fmt"
    "math"
)

func main() {
    arr := []int{3, 2, 4, 1}
    k := 2
    ret := mergeStones1(arr, k)
    fmt.Println(ret)
    fmt.Println("--------")
    ret = mergeStones2(arr, k)
    fmt.Println(ret)
}

func mergeStones1(stones []int, K int) int {
    n := len(stones)
    if (n-1)%(K-1) > 0 {
        return -1
    }
    presum := make([]int, n+1)
    for i := 0; i < n; i++ {
        presum[i+1] = presum[i] + stones[i]
    }
    return process1(0, n-1, 1, stones, K, presum)
}

// part >= 1
// arr[L..R] 一定要弄出part份，返回最低代价
// arr、K、presum（前缀累加和数组，求i..j的累加和，就是O(1)了）
func process1(L int, R int, P int, arr []int, K int, presum []int) int {
    if L == R { // arr[L..R]
        return twoSelectOne(P == 1, 0, -1)
    }
    // L ... R 不只一个数
    if P == 1 {
        next := process1(L, R, K, arr, K, presum)
        if next == -1 {
            return -1
        } else {
            return next + presum[R+1] - presum[L]
        }
    } else { // P > 1
        ans := math.MaxInt64
        // L...mid是第1块，剩下的是part-1块
        for mid := L; mid < R; mid += K - 1 {
            // L..mid(一份) mid+1...R(part - 1)
            next1 := process1(L, mid, 1, arr, K, presum)
            next2 := process1(mid+1, R, P-1, arr, K, presum)
            if next1 != -1 && next2 != -1 {
                ans = getMin(ans, next1+next2)
            }
        }
        return ans
    }
}

func twoSelectOne(c bool, a int, b int) int {
    if c {
        return a
    } else {
        return b
    }
}

func getMin(a int, b int) int {
    if a < b {
        return a
    } else {
        return b
    }
}

func mergeStones2(stones []int, K int) int {
    n := len(stones)
    if (n-1)%(K-1) > 0 { // n个数，到底能不能K个相邻的数合并，最终变成1个数！
        return -1
    }
    presum := make([]int, n+1)
    for i := 0; i < n; i++ {
        presum[i+1] = presum[i] + stones[i]
    }
    dp := make([][][]int, n)
    for i := 0; i < n; i++ {
        dp[i] = make([][]int, n)
        for j := 0; j < n; j++ {
            dp[i][j] = make([]int, K+1)
        }
    }
    return process2(0, n-1, 1, stones, K, presum, dp)
}

func process2(L int, R int, P int, arr []int, K int, presum []int, dp [][][]int) int {
    if dp[L][R][P] != 0 {
        return dp[L][R][P]
    }
    if L == R {
        return twoSelectOne(P == 1, 0, -1)
    }
    if P == 1 {
        next := process2(L, R, K, arr, K, presum, dp)
        if next == -1 {
            dp[L][R][P] = -1
            return -1
        } else {
            dp[L][R][P] = next + presum[R+1] - presum[L]
            return next + presum[R+1] - presum[L]
        }
    } else {
        ans := math.MaxInt64
        // i...mid是第1块，剩下的是part-1块
        for mid := L; mid < R; mid += K - 1 {
            next1 := process2(L, mid, 1, arr, K, presum, dp)
            next2 := process2(mid+1, R, P-1, arr, K, presum, dp)
            if next1 == -1 || next2 == -1 {
                dp[L][R][P] = -1
                return -1
            } else {
                ans = getMin(ans, next1+next2)
            }
        }
        dp[L][R][P] = ans
        return ans
    }
}