3147. Taking Maximum Energy From the Mystic Dungeon

題目

In a mystic dungeon, n magicians are standing in a line. Each magician has an attribute that gives you energy. Some magicians can give you negative energy, which means taking energy from you.

You have been cursed in such a way that after absorbing energy from magician i, you will be instantly transported to magician (i + k). This process will be repeated until you reach the magician where (i + k) does not exist.

In other words, you will choose a starting point and then teleport with k jumps until you reach the end of the magicians' sequence, absorbing all the energy during the journey.

You are given an array energy and an integer k. Return the maximum possible energy you can gain.

Note that when you are reach a magician, you must take energy from them, whether it is negative or positive energy.

題目連結

此題的大意上是給定一個陣列 energy，並且可以任選一個起點，以間隔 k 的方式取得陣列中數字的總和，並且需要找到可以獲得最大總和的值，如以下的例子：

energy = [5, 2, -10, -5, 1]
k = 3

可以有 [5, -5]、[2, 1]、[-10]、[-5]、[1] 這幾種取法，而最大值是由 [2, 1] 得到的總和 3。

思路

一開始直觀的暴力解法是從 index 0 開始，取出間隔 k 的所有元素之後取加總之後的結果，並更新最大值，一路做到最後一項結束，大概會是以下的作法：

def maximumEnergy(self, energy: List[int], k: int) -> int:
    maxEnergy = None
    for idx, _ in enumerate(energy):
        energySum = sum(energy[idx::k])
        maxEnergy = max(energySum, maxEnergy) if maxEnergy is not None else energySum

    return maxEnergy

不過這樣做實際上的複雜度應該會來到 $O(n^2)$，因為在迴圈中的 sum(energy[idx::k] 還是將陣列再走了一遍，只是是以間隔 k 的方式取值。

因此需要觀察每一輪總和之間的關係，從規則中可以發現到如果從第 i 項作為起點的話，那最後的總和會等於從第 i+k 項作為起點的總和再加上第 i 項的值，也就是：

$$energySum_i = energySum_{i+k} + energy[i]$$

所以我們可以使用一個陣列來做記錄，並且從最後一項由後往前來進行，如果這一項往後跳 k 之後超出了陣列的長度，那以該項作為起點的總和就一定會是它自己而已 (像是前面例子的 -10, -5 或 1)：

if (i+k) > len(enengy) - 1:
    sum[i] = energy[i]

那如果間隔 k 之後還有值的話，總和就可以寫成 i+k 項的總和再加上自己：

else:
    sum[i] = sum[i+k] + energy[i]

遍歷一次結束後只要取出 sum 中的最大值即可。

Full Code

class Solution:
    def maximumEnergy(self, energy: List[int], k: int) -> int:
        energyLen = len(energy)
        dp = [0] * energyLen
        for i in range(energyLen - 1, -1, -1):
            if (i+k) > energyLen - 1:
                dp[i] = energy[i]
            else:
                dp[i] = dp[i+k] + energy[i]

        return max(dp)