[정올 1141] 불쾌한 날

문제

소들이 다른 소들의 머리 모양을 얼마나 알 수 있는지 알아본다.

각 소들은 동쪽을 바로보고 있으며 각각 키를 갖고있다. 

소들은 자신보다 키가 작은 소들의 머리만 볼 수 있다.

모든 소들을 대상으로 각각 볼 수 있는 소의 수를 모두 더한 값을 구하라.

 

입력

첫째 줄에 소의 수인 N이 입력된다. (N은 1이상 80,00이하)

이어서 소들의 키가 N개 한 줄로 입력된다.

 

출력

각각의 소들이 볼 수 있는 소의 머리 모양의 총합을 출력한다.

---

접근

N이 최대 80,000인 것을 통해 이중 for문으로 풀 수 없다고 판단하였다. 그래서 한 개의 for문으로 해를 구하는 방법을 생각해보았다.

문제의 예시에서 각 소들이 3+0+1+0+1+0으로 총 5가되는데 첫 번째 소가 볼 수 있는 소들의 마리 수를 어떻게 3마리로 해야할지로 생각하다보니 생각의 발전이 없었다. 

느낌적으로 스택과 같은 자료구조를 활용해야겠다는 생각은 들었지만 어떻게 적용시킬지 생각이 떠오르지 않았다.

 

스택을 활용한다면 다음과 같이 풀 수 있었다.

5 2 4 2 6 1이라는 소들의 키가 입력될 때

처음에 스택이 비어있기 때문에 5를 넣는다.

- 그 다음 2는 스택의 top보다 작기 때문에 스택의 크기만큼 최종해를 증가시키고 2를 스택에 넣는다.

- 그 다음 4는 스택의 top보다 크기 때문에 2를 pop한다.

위와 같은 규칙을 찾아서 해결하였다.

 

구현

- N개의 소에 대한 정보를 입력받으면서 해를 구한다.

- 현재 스택이 비어있다면 스택에 현재 소에 대한 정보를 추가하는 것으로 끝낸다.

- 그렇지 않다면 스택의 top과 비교하여 현재 입력된 소보다 크다면 현재 스택의 길이를 answer에 추가하고 

- 스택에 cow를 추가한다.

- 만약 스택의 top보다 현재 입력한 것이 같거나 크다면 계속해서 스택을 pop한다.

- 스택이 비거나 스택의 top이 큰 것이 나올 때 while문이 종료되면 

- 현재 스택의 길이를 answer에 추가하고 스택에 현재 입력한 소를 추가한다.

for _ in range(N):
    cow = int(input())
    if stack:
        if stack[-1] > cow:
            answer += len(stack)
            stack.append(cow)
        else:
            while stack and stack[-1] <= cow:
                stack.pop()
            answer += len(stack)
            stack.append(cow)
    else:
        stack.append(cow)

---

전체 코드

import sys
input = sys.stdin.readline

N = int(input())
stack = []
answer = 0
for _ in range(N):
    cow = int(input())
    if stack:
        if stack[-1] > cow:
            answer += len(stack)
            stack.append(cow)
        else:
            while stack and stack[-1] <= cow:
                stack.pop()
            answer += len(stack)
            stack.append(cow)
    else:
        stack.append(cow)
print(answer)