문제 개요
숫자로 이루어진 배열 A가 주어졌을 때, i번째 원소를 삭제한 뒤 남은 배열이 좋은 배열(good array)이 되도록 하는 모든 인덱스 i를 찾아야 합니다. 이때 다음 두 가지를 기억해야 합니다.
- 좋은 배열: 배열 안의 어떤 한 원소가 나머지 모든 원소의 합과 정확히 같은 배열
- 인덱스는 1부터 시작하는 1-based indexing을 사용
예시
입력이 [10, 4, 6, 2]라면 출력은 [1, 4]입니다.
- A[1](값 10)을 삭제하면 배열은 [4, 6, 2]가 되고, 6 = 4 + 2이므로 좋은 배열입니다.
- A[4](값 2)를 삭제하면 배열은 [10, 4, 6]이 되고, 10 = 4 + 6이므로 역시 좋은 배열입니다.
접근 방법
핵심 아이디어는 간단합니다. i번째 원소를 삭제한 후 남은 원소들의 총합은 (전체 합 − A[i])입니다. 좋은 배열이 되려면 남은 배열의 어떤 원소 x가 나머지 원소들의 합과 같아야 하므로 다음 식이 성립해야 합니다.
x = (전체 합 − A[i]) − x → x = (전체 합 − A[i]) / 2
따라서 (전체 합 − A[i])가 짝수이고, 그 절반 값이 배열에 실제로 존재하는지만 확인하면 됩니다. 단, A[i]가 그 값과 같은 경우에는 해당 원소를 하나 지우더라도 같은 값이 최소 하나 더 남아 있어야 하므로, 그 값의 빈도가 2 이상인지 추가로 검사해야 합니다.
이를 단계별로 정리하면 다음과 같습니다.
- n := 배열 A의 크기, add := 0(전체 합), my_map := 각 값의 등장 횟수를 저장할 맵
- i를 0부터 n−1까지 반복하며 my_map[A[i]]를 1 증가시키고 add에 A[i]를 더합니다.
- 다시 i를 0부터 n−1까지 반복하며 다음을 수행합니다.
- k := add − A[i]
- k가 짝수이면 k := k / 2로 갱신 (코드의 k >> 1은 비트 연산으로 2로 나누기와 동일)
- k가 my_map에 존재하고, (A[i] == k이면서 my_map[k] > 1) 또는 (A[i] != k)라면 i + 1을 출력
파이썬 구현 예제
다음 구현을 통해 더 잘 이해해 보겠습니다.
from collections import defaultdict
def find_indices(A):
n = len(A)
add = 0
my_map = defaultdict(lambda: 0)
for i in range(n):
my_map[A[i]] += 1
add += A[i]
for i in range(n):
k = add - A[i]
if k % 2 == 0:
k = k >> 1
if k in my_map:
if ((A[i] == k and my_map[k] > 1) or (A[i] != k)):
print(i + 1)
A = [10, 4, 6, 2]
find_indices(A)
입력
[10, 4, 6, 2]
출력
1 4
복잡도 분석
배열을 두 번만 순회하므로 시간 복잡도는 O(n)이며, 각 값의 빈도를 저장하는 맵 때문에 공간 복잡도도 O(n)입니다. 매 삭제마다 배열 전체를 다시 검사하는 O(n²) 완전 탐색보다 훨씬 효율적인 방법입니다.