Chương trình tìm tổng trọng số tối đa cho mảng xoay trong Python

Mô tả bài toán

Cho mảng nums. Tổng trọng số: S = sum(i * nums[i-1]) cho i=1..n (1-indexed). Xoay mảng (cyclic shift) để tối đa hóa S.

Ví dụ

L = [5,3,4]:

• [5,3,4]: 5 + 2*3 + 3*4 = 23
• [3,4,5]: 3 + 2*4 + 3*5 = 26 (max)
• [4,5,3]: 4 + 2*5 + 3*3 = 23

Giải thuật

1. sum_a = sum(nums), cur_val = sum(i * nums[i-1]).
2. Xoay trái lần lượt: cur_val = cur_val - sum_a + nums[i] * n.
3. Cập nhật ans = max(ans, cur_val).

Triển khai Python

def solve(nums):
    n = len(nums)
    sum_a = sum(nums)
    cur_val = ans = sum(nums[i] * (i + 1) for i in range(n))
    for i in range(n):
        cur_val = cur_val - sum_a + nums[i] * n
        ans = max(ans, cur_val)
    return ans

nums = [5, 3, 4]
print(solve(nums))

Đầu vào

[5,3,4]

Đầu ra

26