[LeetCode] No.724 Find pivot Index

문제링크

LeetCode

문제

배열의 인덱스 중에서 인덱스를 기준으로 왼쪽과 오른쪽 원소들의 합이 같은 인덱스를 반환, 합이 같은 경우가 없는 경우, -1반환합니다.

풀이

먼저, 크기가 N인 배열의 해당 인덱스를 기준으로 왼쪽과 오른쪽 원소의 합을 구하면 $O(N)$의 시간복잡도를 가집니다. 배열을 순회하면서 모든 경우를 완전탐색할 경우, 시간복잡도는 최종적으로 $O(N^2)$이 됩니다.

하지만 이를 $O(N)$에 구할 수 있는 방법이 있습니다. 먼저 배열을 순회하면서 원소의 전체 합을 구합니다. 이 때, $O(N)$의 시간복잡도를 가지게 됩니다.
다시 배열을 순회하면서, 오른쪽 합을 전체 합에서 인덱스가 가리키는 배열의 값과 왼쪽합을 빼주어 바로 구합니다. 왼쪽 합과 오른 쪽 값이 같을 때, 바로 인덱스를 반환해주고 아닌 경우, 왼쪽합에 인덱스가 가리키는 배열의 합을 더해줍니다.

배열을 전부 순환한 후에도 조건을 만족하는 인덱스를 찾을 수 없을 경우, -1을 반환해줍니다.

코드

class Solution {
    public int pivotIndex(int[] nums) {

        int totalSum =  Arrays.stream(nums).sum();
        
        int leftSum = 0;
        int rightSum = 0;
        
        for(int i = 0; i<nums.length; i++) {
            
            rightSum=totalSum - nums[i] - leftSum;
            
            if(leftSum == rightSum) {
                return i;
            }
            
            leftSum+=nums[i];
        }
        return -1;
    }
}