문제 소개
처음 N개의 자연수가 한 번씩 등장하는 순열(permutation) 형태의 배열 A와 하나의 정수 M(M ≤ N)이 주어졌을 때, 중앙값이 정확히 M이 되는 연속 부분 배열(sub-array)이 총 몇 개인지 구하는 문제입니다.
중앙값(median)이란 배열의 원소들을 오름차순으로 정렬했을 때 정중앙에 위치한 값을 의미합니다. 단, 길이가 짝수인 경우에는 두 가운데 원소 중 왼쪽에 있는 값을 중앙값으로 사용한다는 점에 유의해야 합니다.
예시
A = [3, 5, 6, 4, 2], M = 5라고 가정해 보겠습니다. 조건을 만족하는 부분 배열은 [3, 5, 6], [5], [5, 6], [5, 6, 4]의 네 가지이므로 정답은 4가 됩니다.
풀이 접근 방법
핵심 아이디어는 각 원소를 M을 기준으로 치환하는 것입니다. M보다 작은 값은 -1, M보다 큰 값은 +1로 바꾸면, 어떤 부분 배열의 중앙값이 M이 되기 위한 조건은 다음과 같이 단순해집니다.
길이가 홀수인 경우: 치환 값들의 합이 0이어야 합니다. 즉, M보다 작은 원소와 큰 원소의 개수가 같아야 합니다.
길이가 짝수인 경우: 왼쪽 가운데 원소가 M이 되어야 하므로 치환 값들의 합이 1이어야 합니다.
따라서 M을 포함하는 부분 배열을 대상으로, 현재까지의 누적 균형값(add)과 시작 지점의 균형값 차이가 0 또는 1이 되는 경우의 수를 세면 됩니다. 이를 위해 M을 만나기 전까지의 균형값을 해시 맵(my_map)에 미리 저장해 두었다가, M 이후부터는 맵 조회만으로 정답을 누적합니다.
알고리즘 단계
- n := 배열 arr의 길이
- my_map := 새 딕셔너리를 만들고 my_map[0] := 1로 초기화 (빈 접두사를 나타냄)
- has := False, add := 0, result := 0으로 초기화
- i를 0부터 n-1까지 반복합니다.
- arr[i] < m이면 add에서 1을 뺍니다.
- arr[i] > m이면 add에 1을 더합니다.
- arr[i] == m이면 has를 True로 설정합니다.
- has가 True라면(M을 이미 지났다면) my_map[add]와 my_map[add - 1]에 해당하는 값이 존재할 때 그 값을 result에 더합니다.
- 그렇지 않다면 my_map[add]의 값을 1 증가시켜 저장합니다.
- 반복이 끝나면 result를 반환합니다.
Python 구현 예제
아래 코드를 통해 실제 동작을 확인해 보겠습니다.
def solve(arr, m):
n = len(arr)
my_map = {}
my_map[0] = 1
has = False
add = 0
result = 0
for i in range(n):
if (arr[i] < m):
add -= 1
elif (arr[i] > m):
add += 1
if (arr[i] == m):
has = True
if (has):
if(add in my_map):
result += my_map[add]
if add-1 in my_map:
result += my_map[add - 1]
else:
my_map[add] = my_map.get(add, 0) + 1
return result
arr = [3, 5, 6, 4, 2]
m = 5
print(solve(arr, m))
실행 결과
입력:
[3, 5, 6, 4, 2], 5
출력:
4
복잡도 분석
배열의 모든 원소를 한 번씩만 순회하므로 시간 복잡도는 O(N)입니다. 해시 맵에는 최대 N개의 균형값이 저장될 수 있으므로 공간 복잡도 역시 O(N)입니다. 덕분에 브루트포스 방식(O(N²))보다 훨씬 효율적으로 큰 입력도 처리할 수 있습니다.