문제 소개
nums라는 배열이 주어졌다고 가정해 봅시다. 우리는 정확히 하나의 인덱스를 골라 그 위치의 요소를 제거할 수 있습니다. 이때 제거된 요소 뒤에 있던 요소들은 한 칸씩 앞으로 당겨지면서 인덱스가 변경됩니다.
배열에서 홀수 인덱스 값들의 합과 짝수 인덱스 값들의 합이 서로 같을 때, 이 배열을 공정한(fair) 배열이라고 정의합니다. 우리가 구해야 할 것은, 하나의 요소를 제거한 뒤 배열이 공정해지도록 만들 수 있는 인덱스의 개수입니다.
예시 확인하기
입력이 nums = [5, 3, 7, 2]일 때 출력은 1입니다. 각 인덱스를 제거했을 때 어떻게 되는지 살펴보겠습니다.
인덱스 0 제거 → [3, 7, 2]: 짝수 위치 합 3+2=5, 홀수 위치 합 7 → 공정하지 않음
인덱스 1 제거 → [5, 7, 2]: 짝수 위치 합 5+2=7, 홀수 위치 합 7 → 공정함
인덱스 2 제거 → [5, 3, 2]: 짝수 위치 합 5+2=7, 홀수 위치 합 3 → 공정하지 않음
인덱스 3 제거 → [5, 3, 7]: 짝수 위치 합 5+7=12, 홀수 위치 합 3 → 공정하지 않음
네 가지 경우 중 조건을 만족하는 것은 인덱스 1을 제거할 때뿐이므로, 정답은 1이 됩니다.
풀이 접근 방법
모든 인덱스마다 실제로 요소를 제거하고 배열을 처음부터 다시 순회하면 O(n²)의 시간이 걸립니다. 하지만 누적합(prefix sum) 아이디어를 활용하면 배열을 두 번만 순회해 O(n) 만에 답을 구할 수 있습니다.
핵심 원리는 다음과 같습니다. 인덱스 i의 요소를 제거하면 i보다 앞의 요소들은 위치가 그대로 유지되지만, i보다 뒤의 요소들은 한 칸씩 앞당겨지면서 홀짝 위치가 서로 바뀌게 됩니다. 따라서 현재 상태의 두 합(sm1, sm2)에서 제거 대상 값만 빼고 새로 포함되는 값을 더해주면, 전체를 다시 계산하지 않고도 다음 경우의 합을 바로 구할 수 있습니다.
알고리즘 단계
- res := 0, sm1 := 0, sm2 := 0으로 초기화합니다.
- 인덱스 1부터 배열 끝까지 순회하며, i가 홀수면 sm1에 nums[i]를 더하고, i가 짝수면 sm2에 nums[i]를 더합니다. (인덱스 0은 제외)
- 이 상태는 '인덱스 0을 제거한 경우'와 같으므로, sm1 == sm2이면 res를 1 증가시킵니다.
- 다시 인덱스 1부터 끝까지 순회하며 제거 대상을 하나씩 옮겨 갑니다. i가 홀수면 sm1 = sm1 − nums[i] + nums[i−1], i가 짝수면 sm2 = sm2 − nums[i] + nums[i−1]로 갱신합니다.
- 갱신 후 sm1 == sm2이면 res를 1 증가시킵니다.
- 모든 순회가 끝나면 res를 반환합니다.
파이썬 구현 코드
def solve(nums): res, sm1, sm2 = 0, 0, 0 # 인덱스 0을 제거한 경우에 해당하는 초기 합 계산 for i in range(1, len(nums)): if i % 2 == 1: sm1 += nums[i] else: sm2 += nums[i] if sm1 == sm2: res += 1 # 인덱스 1부터 차례대로 제거하는 경우로 상태 갱신 for i in range(1, len(nums)): if i % 2 == 1: sm1 = sm1 - nums[i] + nums[i-1] else: sm2 = sm2 - nums[i] + nums[i-1] if sm1 == sm2: res += 1 return res nums = [5, 3, 7, 2] print(solve(nums))
실행 결과
입력:
[5, 3, 7, 2]
출력:
1
복잡도 분석
- 시간 복잡도: O(n) — 배열을 총 두 번 순회합니다.
- 공간 복잡도: O(1) — 추가 메모리 없이 변수 몇 개만 사용합니다.