724. Find Pivot Index

양성준·2025년 4월 8일

코딩테스트

목록 보기
15/102

문제

https://leetcode.com/problems/find-pivot-index/description/

풀이

1. 처음 풀이 (prefixSum 배열 이용)

class Solution {
    public int pivotIndex(int[] nums) {
        int[] prefixSum = new int[nums.length];
        int[] suffixSum = new int[nums.length];

        prefixSum[0] = nums[0];
        suffixSum[nums.length-1] = nums[nums.length-1];

        // prefixSum 먼저 구하기
        for(int i = 1; i < nums.length; i++) {
            prefixSum[i] = prefixSum[i-1] + nums[i];
        }

        // suffixSum 구하기
        for(int i = nums.length - 2; i >= 0; i--) {
            suffixSum[i] = suffixSum[i+1] + nums[i];
        }

        // 중요한건 어떻게든 O(N)으로 푸는 것
        // i를 기준으로, leftSum과 rightSum이 같다면 그것이 pivot index인것
        // i를 이동시켜가면서 왼쪽과 오른쪽의 합을 비교
         for (int i = 0; i < nums.length; i++) {
            int leftSum = i == 0 ? 0 : prefixSum[i - 1]; // 맨 왼쪽인 경우 0, 아니라면 i의 한칸 왼쪽
            int rightSum = i == nums.length - 1 ? 0 : suffixSum[i + 1]; // 맨 오른쪽인 경우 0, 아니라면 i의 한칸 오른쪽
            
            if (leftSum == rightSum) return i;
        }
        return -1;
    }
}

  • 핵심은 i를 기준으로 prefixSum과 suffixSum이 같아야 Pivot index인 것!
  • O(N)의 시간복잡도지만, 별도의 배열을 두개나 써서 공간복잡도가 높음

2. 다른 풀이 (반복문 내부에서 leftSum구하고, 그걸로 rightSum 구하기)

class Solution {
    public int pivotIndex(int[] nums) {
        int leftSum = 0;
        int totalSum = 0;

        for(int i = 0; i < nums.length; i++) {
            totalSum += nums[i];
        }

        // 중요한건 i를 기준으로 leftSum과 rightSum이 같아야함! 
         for (int i = 0; i < nums.length; i++) {
            // i를 기준으로 한칸왼쪽 까지의 left Sum과
            leftSum = i == 0 ? 0 : leftSum+nums[i-1];
            // 한칸오른쪽 까지의 right Sum을 구해줌
            int rightSum = i == nums.length-1 ? 0 : totalSum - leftSum - nums[i];

            if(leftSum == rightSum) {
                return i;
            }
        }
        return -1;
    }
}

  • rightSum은 totalSum에서 leftSum과 nums[i]를 뺀 것
  • 코드가 훨씬 간편해짐! + 공간복잡도 O(1)의 풀이
profile
백엔드 개발자

0개의 댓글