파이썬에서 리스트의 최대 파워(Power) 값 구하는 프로그램

리스트의 파워(power)는 모든 인덱스에 대해 (인덱스 + 1) × 해당 위치의 값을 곱한 뒤 모두 더한 합으로 정의됩니다. 수식으로 나타내면 다음과 같습니다.

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

문제 설명

N개의 양의 정수로 이루어진 리스트 nums가 주어졌다고 가정해 보겠습니다. 우리는 리스트에서 임의의 값을 하나 선택해 교환(swap)이 아닌 이동(move) 방식으로 원하는 위치, 즉 리스트의 맨 앞이나 맨 끝을 포함한 어느 곳으로든 옮길 수 있습니다. 물론 아무 요소도 옮기지 않는 것도 허용됩니다. 이렇게 조작했을 때 리스트가 가질 수 있는 최대 파워 값을 구하고, 그 결과를 10^9 + 7로 나눈 나머지를 반환해야 합니다.

예를 들어 입력이 nums = [4, 2, 1]이라면 출력은 16이 됩니다.

풀이 접근 방법

이 문제는 누적 합(prefix sum)과 볼록 껍질 트릭(convex hull trick)을 활용하면 효율적으로 해결할 수 있습니다. 핵심 단계는 다음과 같습니다.

P := [0] 으로 누적 합 배열을 초기화합니다.
base := 0 으로 초기 파워 값을 설정합니다.
각 인덱스 i와 원소 x에 대해 다음을 반복합니다.
- P의 끝에 P[-1] + x를 추가합니다.
- base := base + i * x 로 갱신합니다.
eval_at(j, x) 함수를 정의합니다. 이 함수는 -j * x + P[j]를 반환합니다.
intersection(j1, j2) 함수를 정의합니다. 이 함수는 두 직선의 교차점을 계산하여 (P[j2] - P[j1]) / (j2 - j1)을 반환합니다.
hull := [-1], indexes := [0] 으로 초기화합니다.
j를 1부터 P의 크기까지 순회하며 볼록 껍질을 구성합니다.
- hull이 비어 있지 않고 intersection(indexes[-1], j) <= hull[-1]인 동안 hull과 indexes의 마지막 원소를 제거합니다.
- hull 끝에 intersection(indexes[-1], j)를 추가하고, indexes 끝에 j를 추가합니다.
ans := base 로 초기화합니다.
각 인덱스 i와 원소 x에 대해 다음을 수행합니다.
- x가 정렬 순서를 유지하며 hull에 삽입될 수 있는 위치 j를 이분 탐색(bisect)으로 찾습니다.
- j := max(j - 1, 0) 으로 조정합니다.
- ans := max(ans, base + eval_at(i, x) - eval_at(indexes[j], x)) 로 갱신합니다.
ans mod (10^9 + 7)을 반환합니다.

구현 예시

아래 코드를 통해 더 자세히 이해할 수 있습니다.

import bisect
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
        def eval_at(j, x):
            return -j * x + P[j]
        def intersection(j1, j2):
            return (P[j2] - P[j1]) / (j2 - j1)
        hull = [-1]
        indexes = [0]
        for j in range(1, len(P)):
            while hull and intersection(indexes[-1], j) <= hull[-1]:
                hull.pop()
                indexes.pop()
            hull.append(intersection(indexes[-1], j))
            indexes.append(j)
        ans = base
        for i, x in enumerate(A):
            j = bisect.bisect(hull, x)
            j = max(j - 1, 0)
            ans = max(ans, base + eval_at(i, x) - eval_at(indexes[j], x))
        return ans % (10 ** 9 + 7)

ob = Solution()
print (ob.solve([4, 2, 1]))

입력

[4, 2, 1]