LeetCode/NeetCode

[Prefix Sums] 724. Find Pivot Index

hyunkookim 2025. 4. 5. 07:26

724. Find Pivot Index

 

1) Prefix Sums 누적합 이용 풀이 방법

class Solution:
    def pivotIndex(self, nums: List[int]) -> int:
        # 자신을 제외한 왼쪽 합 == 오른쪽 합이 되는 pivot index를 찾는 문제

        total = 0                     # 전체 합을 계산하기 위한 변수
        PrefixSum = []               # 0번 인덱스부터의 누적합 배열

        for n in nums:
            total += n               # 현재까지의 누적합을 계산
            PrefixSum.append(total)  # 누적합을 배열에 저장

        # 구간 합을 구하는 헬퍼 함수: nums[s]부터 nums[e]까지의 합 반환
        def getSum(s, e):
            if s > e:
                return 0             # 구간이 유효하지 않을 경우 합은 0
            l = PrefixSum[s-1] if s > 0 else 0  # 시작 지점 바로 앞까지의 합
            r = PrefixSum[e]                    # 끝 지점까지의 누적합
            return r - l                        # 구간 [s, e]의 합

        # 모든 인덱스를 돌면서 pivot 조건이 맞는지 확인
        for i in range(len(nums)):
            leftSum = getSum(0, i-1)                 # i 왼쪽 구간의 합
            RightSum = getSum(i+1, len(nums)-1)      # i 오른쪽 구간의 합
            if leftSum == RightSum:                  # 왼쪽합 == 오른쪽합이면
                return i                              # 현재 인덱스가 pivot
        return -1                                     # 조건을 만족하는 인덱스가 없으면 -1 반환

 

2) prefix sum 배열 없이 공간을 절약하면서 pivot index를 찾는 최적화된 풀이

class Solution:
    def pivotIndex(self, nums: List[int]) -> int:
        total = sum(nums)        # 전체 합
        leftSum = 0              # 왼쪽 누적합

        for i, n in enumerate(nums):
            # 왼쪽 합 == 오른쪽 합 → total - left - n == left
            # → 즉, total - n == 2 * leftSum
            if leftSum == total - leftSum - n:
                return i
            leftSum += n
        return -1

✅ 어떻게 작동하냐면?

 
왼쪽 합 == 오른쪽 합
if leftSum == total - leftSum - n: 
    return i
  • total - leftSum - n → # 현재 원소 n을 제외한 오른쪽 합

✅ 장점

  • 시간 복잡도: O(n) (한 번만 순회)
  • 공간 복잡도: O(1) (prefix 배열 없이)