门| GATE-CS-2005 |问题 10

本题涉及编程和算法。

题目描述

给定一个长度为 $n$ 的整数序列 $a_1, a_2, \ldots, a_n$，你需要设计一个算法，找到一个子序列 $a_{i_1}, a_{i_2}, \ldots, a_{i_k}$ 使得该子序列的和最大，并且 $a_{i_j} - a_{i_{j-1}}$ 不等于 $d$，其中 $d$ 是所给的一个整数。

实现一个函数 find_max_sum(arr: List[int], d: int) -> Tuple[int, List[int]]，该函数输入为整数序列 arr 和整数 d，输出为一个元组，第一个元素表示所找到子序列的最大和，第二个元素为所找到的子序列本身。

数据规模

$1 \leq n \leq 10^4$，$-10^3 \leq a_i \leq 10^3$，$-10^3 \leq d \leq 10^3$。

算法思路

本题可以使用动态规划求解。我们定义 $dp_i$ 表示以 $a_i$ 作为子序列结尾的最大和，并且该子序列中相邻的元素之差不等于 $d$。则有：

$$ dp_i = \max\limits_{j<i, a_i-a_j\neq d}{dp_j} + a_i $$

初始状态为 $dp_1 = a_1$。最终的答案为 $\max\limits_{i=1}^n{dp_i}$。

基于上述状态转移方程，我们可以快速地求出最大和。但我们需要进一步记录下哪些元素形成了最终的答案。为此，我们在状态转移时记录下状态 $dp_i$ 对应最优解在哪个状态 $dp_j$ 上转移而来，即：

$$ prev_i = \arg\max\limits_{j<i, a_i-a_j\neq d}{dp_j} $$

然后可以通过倒序遍历这个数组，依次记录下所有 $\arg\max$ 状态的元素，即为最终结果。

代码实现

from typing import List, Tuple

def find_max_sum(arr: List[int], d: int) -> Tuple[int, List[int]]:
    n = len(arr)
    dp = arr.copy()
    prev = [-1] * n

    for i in range(1, n):
        for j in range(i):
            if arr[i] - arr[j] == d:
                continue
            tmp = dp[j] + arr[i]
            if tmp > dp[i]:
                dp[i] = tmp
                prev[i] = j

    max_idx = max(enumerate(dp), key=lambda x: x[1])[0]

    res = []
    idx = max_idx
    while idx != -1:
        res.append(arr[idx])
        idx = prev[idx]

    res.reverse()

    return dp[max_idx], res

参考资料

• GeeksforGeeks - Maximum subsequence sum such that no two elements in the subsequence are adjacent
• LeetCode - 740. Delete and Earn（一道与本题类似的 LeetCode 题目，进一步阐释了本题的动态规划思想）