[BaekJoon] 3015 오아시스 재결합 (Platinum 5) - Python

문제

 

오아시스의 재결합 공연에 N명이 한 줄로 서서 기다리고 있다.

이 역사적인 순간을 맞이하기 위해 줄에서서 기다리고 있던 백준이는 갑자기 자기가 볼 수 있는 사람의 수가 궁금해 졌다.

두 사람 A와 B가 서로 볼 수 있으려면, 두 사람 사이에 A 또는 B보다 키가 큰 사람이 없어야 한다.

줄에 서있는 사람의 키가 주어졌을 때, 서로 볼 수 있는 쌍의 수를 구하는 프로그램을 작성하시오.

---

입출력 예시

 

입력: 첫째 줄에 줄에서 기다리고 있는 사람의 수 N이 주어진다. (1 ≤ N ≤ 500,000)

둘째 줄부터 N개의 줄에는 각 사람의 키가 나노미터 단위로 주어진다. 모든 사람의 키는 231 나노미터 보다 작다.

사람들이 서 있는 순서대로 입력이 주어진다.

출력:

서로 볼 수 있는 쌍의 수를 출력한다.

 

입력예시1)

7
2
4
1
2
2
5
1

출력예시1)

10

---

코드

 

이 문제는,,, 스택을 2차원으로 이용해야 풀 수 있다. (대충 (x, y)의 형태로 append, pop해야한다는 뜻)

 

1. 현재보다 키가 더 작은 경우

2. 현재와 키가 같은 경우

3. 현재보다 키가 더 큰 경우

 

이렇게 3가지로 나누어서 구성한다.

 

1. 현재보다 키가 더 작은 경우

- 이 경우, while문을 이용하여 stack의 맨 위 값이 현재 값보다 작을 때만 확인한다. stack에서 pop을 하면 (h, buf)형태로 나오고, 뒤의 buf는 h와 키가 같은 사람이 연속으로 서있는 수를 말한다. 그래서 결과값에 buf를 하나씩 더해주면 된다. 만약 h값이 현재보다 작은 데이터를 만나면 반복을 멈춘다.

 

2. 현재와 키가 같은 경우

- 이 경우 스택의 맨 위 데이터를 pop한다. 이때 나오는 buf값을 결과값에 더해주고, pop한 이후에 스택에 값이 남아있다면 결과값에 1을 더 더해준다. 아니라면 더하지 않는다. 마지막으로 스택에 (현재키, buf + 1)을 추가한다.

 

3. 키가 더 큰 경우

- 키가 더 큰 경우는 그냥 스택 맨 위에 (현재키, 1)을 추가하고 결과값을 1 증가시킨다.

 

 

import sys
input = sys.stdin.readline
N = int(input()) # N입력
num = [int(input()) for _ in range(N)] # N만큼 반복하여 리스트에 숫자 입력
stack = [] # 문제를 풀기 위한 스택 선언
cnt = 0 # 결과값 선언

for i in num: # 리스트의 값을 하나씩 가져온다.
    while stack and stack[-1][0] < i: # 만약 스택에 값이 들어있고, 현재의 값이 스택의 마지막 값의 키 정보보다 작다면
        cnt += stack.pop()[1] # 결과값에 스택 맨 위의 값을 pop하면서, 키가 같은 사람의 수를 더해준다
    if not stack: # 만약 스택이 비어있다면
        stack.append((i, 1)) # 현재의 값을 스택에 추가하고 continue한다.
        continue
    if stack[-1][0] == i: # 만약 스택이 비어있지 않고 스택 맨 위의 값이 현재의 값과 같다면
        buf = stack.pop()[1] # 스택 맨 위의 키가 같은 사람의 수 값을 buf에 저장하고
        cnt += buf # 결과값을 buf만큼 추가한다.
        if stack: # 만약 위의 pop이후에 스택이 비어있지 않다면
            cnt += 1 # 결과값을 1 증가시킨다.
        stack.append((i, buf + 1)) # 스택의 맨 위에 현재의 값과 키가 같은 사람의 수를 1 증가하여 추가한다.
    else: # 만약 스택이 비어있지 않고, 스택의 맨 위의 값이 같거나 작지 않다면,
        stack.append((i, 1)) # 스택의 맨 위에 현재의 값 i와 키가 같은 사람의 수(초기값 1)을 추가한다.
        cnt += 1 # 결과값에 1 추가한다.
    
print(cnt) # 결과값 출력

---

실행 화면

 

---

채점 결과

input함수 재정의를 안해서 그런가 시간초과가 나드라 ㅇ.ㅇ