Computer >> Computer tutorials >  >> Programming >> Python

Program to find maximum weighted sum for rotated array in Python

Suppose we have an array of few elements. We shall have to find the maximum weighted sum if the array elements are rotated. The weighted sum of an array nums can be calculated like below −

$$\mathrm{𝑆=\sum_{\substack{𝑖=1}}^{n}𝑖∗𝑛𝑢𝑚𝑠[𝑖]}$$

So, if the input is like L = [5,3,4], then the output will be 26 because

  • array is [5,3,4], sum is 5 + 2*3 + 3*4 = 5 + 6 + 12 = 23

  • array is [3,4,5], sum is 3 + 2*4 + 3*5 = 3 + 8 + 15 = 26 (maximum)

  • array is [4,5,3], sum is 4 + 2*5 + 3*3 = 4 + 10 + 9 = 23

To solve this, we will follow these steps −

  • n := size of nums
  • sum_a := sum of all elements in nums
  • ans := sum of all elements of (nums[i] *(i + 1)) for all i in range 0 to n
  • cur_val := ans
  • for i in range 0 to n - 1, do
    • cur_val := cur_val - sum_a + nums[i] * n
    • ans := maximum of ans and cur_val
  • return ans

Example

Let us see the following implementation to get better understanding −

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))

Input

[5,3,4]

Output

26