파이썬 오아시스 재결합(BOJ 3015)


문제

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

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

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

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


입력

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

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

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


출력

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


예제 입력 1

7
2
4
1
2
2
5
1

예제 출력 1

10


📝 풀어보기

스택을 만들어서 (높이, 카운트)를 같이 저장하고 조건에 따라서 스택을 하나씩 비우면서, 저장해둔 카운트 값을 ans에 누적시켰다.

스택에 값이 들어있을때, 스택이 비어있고 끝자리가 높이보다 작아질때까지 스택을 비우면서 ans에 카운트 값을 누적한다.

스택의 마지막 값과 들어온 높이가 같으면 서로 볼 수 있는 상태이므로 스택을 비우면서 카운트를 누적한다. 그리고 입력한 (높이, 카운트+1)을 스택에 누적시킨다.

이 때, 스택의 길이가 0이 아니라면 ans에 값을 1 누적한다.

스택이 비어있는 경우엔 길이가 0이 아니라면 ans에 값을 1 누적하고 스택에 (높이, 1)을 저장한다.

import sys
input = sys.stdin.readline

stack = []
ans = 0

# N명이 한 줄로 서있다.
for _ in range(int(input())):
    hight = int(input())

    # 스택이 비어있고, 끝자리가 높이보다 작으면 좋료
    while stack and stack[-1][0] < hight:
        # print(stack)
        # 카운트값을 꺼내 ans에 누적한다
        ans += stack.pop()[-1]
        # print(stack)

    # 스택의 끝자리가 높이와 같을때
    if stack and stack[-1][0] == hight:
        # 카운트를 누적한다.
        cnt = stack.pop()[1]
        ans += cnt

        # 스택의 길이가 0이 아니라면? 스택에 수가 들어왔을때 
        if len(stack) != 0:
            ans += 1
        # 스택의 끝자리와 높이가 같으면 (높이, 카운트+1) 값을 추가
        stack.append((hight, cnt+1))
        # print(stack)
    # 스택이 비어있는 경우
    else:
        if len(stack) != 0:
            ans += 1
        stack.append((hight, 1))   

print(ans)

관심있을 포스팅