环形数组中的最大贡献值 | 豆包MarsCode AI刷题

问题描述

小S拿到了一个长度为 nn 的环形数组，并定义了两个下标 ii 和 jj 的贡献值公式为：
f(i, j) = (a_i + a_j) × dist(i, j)
其中 dist(i, j) 是下标 ii 和 jj 在数组中的最短距离。小S希望找到一对下标，使得它们的贡献值尽可能大。环形数组的特点是最左和最右的元素也是相邻的。你需要帮助她找到最大贡献值。

例如，给定数组 [1, 2, 3]，由于是环形数组，任意两个下标的距离都是1，因此 f(2,3)=(2+3)×1=5f(2,3)=(2+3)×1=5。

---

测试样例

样例1：

输入：n = 3,a = [1, 2, 3]
输出：5

样例2：

输入：n = 4,a = [4, 1, 2, 3]
输出：12

样例3：

输入：n = 5,a = [1, 5, 3, 7, 2]
输出：24

---

问题理解

你有一个长度为 n 的环形数组 a，需要找到两个下标 i 和 j，使得它们的贡献值 f(i, j) = (a[i] + a[j]) * dist(i, j) 最大。其中 dist(i, j) 是下标 i 和 j 在数组中的最短距离。

数据结构选择

由于数组是环形的，最短距离 dist(i, j) 可以通过以下公式计算：

• dist(i, j) = min((j - i + n) % n, (i - j + n) % n)

算法步骤

1. 初始化最大贡献值：用一个变量 result 来存储当前找到的最大贡献值，初始值为 0。

2. 双重循环遍历所有下标对：

   • 外层循环遍历下标 i，内层循环遍历下标 j。
   • 确保 i 和 j 不相等，因为同一个下标的贡献值为 0。

3. 计算最短距离：

   • 使用公式 dist(i, j) = min((j - i + n) % n, (i - j + n) % n) 计算 i 和 j 的最短距离。

4. 计算贡献值：

   • 计算 contribution = (a[i] + a[j]) * dist(i, j)。

5. 更新最大贡献值：

   • 如果当前的 contribution 大于 result，则更新 result。

6. 返回结果：

   • 循环结束后，result 即为最大贡献值。

---

代码解释

• 双重循环：遍历所有可能的下标对 (i, j)。
• 最短距离计算：使用 min((j - i + n) % n, (i - j + n) % n) 计算最短距离。
• 贡献值计算：计算 (a[i] + a[j]) * dis 并更新最大值。
• 测试用例：main 函数中包含了三个测试用例，分别测试了不同长度的数组，并验证了结果是否正确。

完整代码

def solution(n, a):
    result = 0
    for i in range(n):
        for j in range(n):
            if i != j:
                dis = min((j - i + n) % n, (i - j + n) % n)
                contribution = (a[i] + a[j]) * dis
                result = max(result, contribution)
    return result

def main():
    print(solution(3, [1, 2, 3]) == 5)
    print(solution(4, [4, 1, 2, 3]) == 12)
    print(solution(5, [1, 5, 3, 7, 2]) == 24)

if __name__ == "__main__":
    main()