문제 출처 : https://www.acmicpc.net/problem/3015
문제 파악
N명이 한 줄로 서 있고, 서로 볼 수 있는 쌍의 수를 구하는 문제.
예시 입력
7 # 사람 수
2
4
1
2
2
5
1
정답은 10. 직관적으로 세어보면,
- 서로 바로 인접한 사람들은 무조건 서로 볼 수 있으므로
(N - 1)쌍이 기본. - 그 외에도, A와 B 사이에 있는 사람들 중 A, B 보다 ‘둘 다 크지 않으면’(작거나 같으면) A와 B는 서로 볼 수 있음.
내가 처음 손으로 센 쌍들:
2-4, 4-1, 1-2, 2-2, 2-5, 5-1 → 6쌍
여기에,
4-(1)-2, 4-(1,2)-2, 4-(1,2,2)-5, 2-(2)-5 → 4쌍
합쳐서 10쌍.
N ≤ 500,000 이라서 O(N^2) 는 절대 안 됨.
한 번의 스캔으로 끝내는 O(N) 풀이가 필요하고, 모노톤(단조) 스택이라는 방법을 써야 한다.
핵심 아이디어
스택에는 현재 후보보다 앞에 있는 사람 중에서 아직 뒤 사람이랑 연결될 가능성이 남아 있는 후보들의 (height, count) 를 넣는다.
여기서 count 는 “같은 키(height)가 연속으로 몇 명 있는지”의 개수다. (예: ..., (2,3)이라면 키 2가 연속 3명)
새 사람의 키 h를 처리할 때 규칙:
1) 스택 top의 키가 h보다 작은 경우
top.height < h인 동안 pop 하면서 그 그룹의 count만큼 쌍을 더한다.- 이유: 더 작은 사람들은
h와 서로 볼 수 있고,h가 오면서 가려지므로 이제 기회가 끝난다.
2) 스택 top의 키가 h와 같은 경우
- 같은 키의 그룹
c를 pop 하고, 정답 += c (서로 다 보임) - 그리고 현재 같은 키 그룹을 합쳐서
count += c - 스택에 아직 누군가 남아 있다면(= 더 큰 키가 뒤에 존재), 그 더 큰 사람과도 한 쌍이 더 생김 → 정답 += 1
(같은 키 그룹 뒤에 더 큰 키가 “경계”처럼 하나 더 보이는 구조)
3) 스택 top의 키가 h보다 큰 경우
- pop 하지 않고, 정답 += 1 (바로 앞 큰 사람 하나는 어쨌든 보임)
마지막에 (h, count) 를 스택에 push.
이렇게 하면 각 사람(각 노드)은 최대 한 번 push/pop 되므로 총 O(N) 으로 풀 수 있다.
해답 및 풀이
import sys
input = sys.stdin.readline
N = int(input().strip()) # 사람 수
stack = [] # (height 키 , count 같은 키 연속 개수) 형태로 스택에 저장 할 것임
ans = 0 # 서로 볼 수 있는 쌍의 개수
for _ in range(N): # N명의 사람을 왼쪽에서 오른쪽으로 순차 처리할 것임
h = int(input().strip()) # 키 입력받음
cnt = 1 # 현재 키 h를 가진 사람의 연속 개수, 초기값 1명
# 1) 스택 top의 키가 이제 들어오는 현재 키보다 작은 경우,
while stack and stack[-1][0] < h:
ans += stack[-1][1] # # 해당 stack의 그룹들의 인원 수 만큼 ans에 쌍 개수를 추가
stack.pop() # 가려져서 더 이상 뒤와 연결될 수 없으므로 제거
if stack and stack[-1][0] == h: # 같은 키 그룹이 바로 뒤에 오는 경우
c = stack[-1][1]
ans += c # 같은 키 끼리는 전부 볼 수 있으므로 c만큼 쌍 추가
stack.pop() # 기존 같은 키 그룹은 병합할 것이므로 제거
cnt += c # 현재 사람까지 포함해 같은 키 연속 개수 갱신
if stack: # 같은 키 그룹 뒤에 더 큰 키가 남아 있다면
ans += 1 # 그 더 큰 사람 1명과도 서로 볼 수 있으므로 +1
else:
if stack:# 스택이 비었거나, top의 키가 이제 오는 현재 키 보다 더 큰 경우
ans += 1 # 바로 앞의 더 큰 사람 1명과는 서로 볼 수 있으므로 +1
stack.append((h,cnt))
print(ans)
예시 흐름 (2, 4, 1, 2, 2, 5, 1)
스택 표기: [(키, 개수)], ans는 누적 쌍 수
1) h=2
- 스택 비어있음 → push (2,1)
- ans=0
2) h=4
- top(2)<4 → pop하며 ans +=1 → ans=1
- 스택 비움 → push (4,1)
- ans=1
3) h=1
- top(4)>1 → ans +=1 → ans=2
- push (1,1)
- 스택: [(4,1),(1,1)]
4) h=2
- top(1)<2 → pop, ans +=1 → ans=3
- 이제 top(4)>2 → ans +=1 → ans=4
- push (2,1)
5) h=2
- top(2)==2 → 같은 그룹 c=1 pop, ans +=1 → ans=5
- count=1+1=2
- 스택에는 (4,1) 남아있음 → ans +=1 → ans=6
- push (2,2)
6) h=5
- top(2)<5 → pop, ans +=2 → ans=8
- top(4)<5 → pop, ans +=1 → ans=9
- 스택 비움 → push (5,1)
7) h=1
- top(5)>1 → ans +=1 → ans=10
- push (1,1)
최종 ans = 10
stack에 튜플형식으로 저장하며 푸는 단조스택 문제였다.
모노톤(단조) 스택에 대해 정리하고 다시 봐야겠다.
