Maximum value of K coins taken from P piles in Go, using Dynamic Programming

k_from_p.go

// Leetcode:
//    https://leetcode.com/problems/maximum-value-of-k-coins-from-piles/

func maxValueOfCoins(piles [][]int, k int) int {
    // Record the points cumulatively for taking N coins from each pile
    for p := range piles {
        for c := range piles[p] {
            if c+1 >= len(piles[p]) {
                continue
            }
            piles[p][c+1] += piles[p][c]
        }
    }
    // Now, piles[p][c] is the value of taking c coins from pile p
    
    
    // We're going to make a 2d state array, in which
    //   state[c][p]
    // contains the maximum value of spending c coins using the first p piles.
    // This requires k+1 * len(piles)+1 to store, because each dimension is inclusive.
    state := make2d(k + 1, len(piles)+1)
    //
    // Our goal is to compute state[k][len(piles)],
    // which will contain the max value of spending all coins using all piles.
    
    // Compute state[*][p+1] based on taking coins from piles[p].
    for p := 0; p < len(piles); p++ {
        // Compute the best for c or more coins based on c coins and taking more from p.
        for c := 0; c <= k; c++ {
            current := state[c][p]
            // Given the best way to spend c coins using p piles (i.e. state[c][p]),
            // We can compute new ways of spending c+n coins by:
            //   taking n coins from pile p (i.e. state[c][p] + piles[p][n-1])
            //
            // First, the n = 0 case: if we don't take any coins from this pile,
            //   it'll still have the same value when we come around to the next pile.
            // Unless we already found a better way to get there, obviously.
            if current > state[c][p+1] {
                state[c][p+1] = current
            }
            // Then we can consider taking nonzero numbers of coins.
            //   There's no need to check spending too many coins (c+n > k)
            //     or spend more coins than the pile (piles[p]) contains.
            for n := 1; c+n <= k && n-1 < len(piles[p]); n++ {
                // Note that pile[n] contains the cumulative value of taking n+1 coins, since 0 coins is assumed to be 0,
                //   so the value of taking n coins from pile p is piles[p][n-1]
                nFromPile := piles[p][n-1]
                value := current + nFromPile
                totalCoins := c+n
                // Only update the value if we've found a better way to spend the coins.
                if value > state[totalCoins][p+1] {
                    state[totalCoins][p+1] = value
                }
            }
        }
    }
    
    // To get here, we had to go through:
    //    k coins
    //  * p piles
    //  * k ways of spending coins from the pile
    //  = O(k^2 * p)
    return state[k][len(piles)]
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

func maxs(s []int) (max int) {
    for _, v := range s {
        if v > max {
            max = v
        }
    }
    return
}

func make2d(rows, cols int) [][]int {
    back := make([]int, rows*cols)
    out := make([][]int, rows)
    for i := range out {
        out[i] = back[i*cols:][:cols:cols]
    }
    return out
}