Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬에서 요소 재배치로 리스트의 파워(Power) 최댓값 구하는 프로그램

문제 개요

N개의 양수로 이루어진 리스트 nums가 있다고 가정해 보겠습니다. 이때 우리는 리스트에서 임의의 값을 하나 골라 다른 위치로 이동(교환이 아니라 이동)할 수 있으며, 아예 아무것도 이동하지 않아도 됩니다. 목표는 리스트의 최종 파워(power)가 가장 커지도록 만드는 것입니다.

여기서 리스트의 파워란 모든 인덱스 i에 대해 (인덱스 + 1) × 해당 위치의 값의 합으로 정의됩니다.

$$\displaystyle\sum\limits_{i=0}^{n-1} (i+1)\times list[i]$$

예를 들어 입력이 nums = [6, 2, 3]이라면 결과는 26이 됩니다. 값 6을 맨 뒤로 옮겨 리스트를 [2, 3, 6]으로 만들면 파워가 (2 × 1) + (3 × 2) + (6 × 3) = 26이 되기 때문입니다.

해결 접근 방법

이 문제는 누적합(prefix sum)을 활용하면 효율적으로 해결할 수 있습니다. 요소를 한 위치에서 다른 위치로 이동할 때 파워의 변화량을 수식으로 계산하면, 실제로 리스트를 재배열하지 않고도 최댓값을 구할 수 있습니다. 구체적인 단계는 다음과 같습니다.

  • P := 값이 0인 리스트로 초기화 (누적합 저장용)

  • base := 0 (원래 리스트의 초기 파워)

  • A의 각 인덱스 i와 값 x에 대해 다음을 수행:

    • P의 마지막 원소에 x를 더한 값을 P의 끝에 추가

    • base := base + (i+1) × x

  • ans := base

  • A의 각 인덱스 i와 값 x에 대해 다음을 수행:

    • j를 0부터 len(A)까지 반복:

      • ans := max(ans, base + P[i] − P[j] − (i − j) × x)

  • ans 반환

구현 예제

더 잘 이해하기 위해 다음 파이썬 구현 코드를 살펴보겠습니다.

class Solution:
   def solve(self, A):
      P = [0]
      base = 0
      for i, x in enumerate(A, 1):
         P.append(P[-1] + x)
         base += i * x
      ans = base
      for i, x in enumerate(A):
         for j in range(len(A) + 1):
            ans = max(ans, base + P[i] - P[j] - (i - j) * x)
      return ans
ob = Solution()
nums = [6, 2, 3]
print(ob.solve(nums))

입력

[6, 2, 3]

출력

26

정리

이 알고리즘은 먼저 원본 리스트의 파워(base)와 누적합 배열(P)을 계산한 뒤, 각 요소를 가능한 모든 위치로 이동시켰을 때의 파워 변화를 수식으로 평가합니다. 시간 복잡도는 O(N²)이며, 리스트를 실제로 조작하지 않고도 최적의 이동 전략을 찾을 수 있다는 점이 핵심입니다.