[Java+Python]3015번, 오아시스 재결합

목차

 

문제

 

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

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

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

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

 

입력

 

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

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

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

 

출력

 

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

 

풀이

 

솔직히 이거 혼자서 풀어보려고 며칠을 고생했다.

 

스택 문제가 생각보다 만만해서(?) 그런지 뭔가 혼자서 할 수 있을 것 같았는데..

 

결국 어림도 없어서 다른 사람의 풀이와 아이디어를 참고했다.

 

애초에 생각을 다르게 해야 하는 부분이 있었는데, 일단 설명해 보겠다.

 

문제를 요약하면 "키가 큰 사람에게 가리지 않고 볼 수 있는 사람의 쌍"을 찾는 것인데,

 

이를 해결하기 위해서는 '키'와 '키가 같은 사람의 수'를 필드로 갖는 객체를 만들어서 스택에 넣어야 했다.

 

이걸 받아들이고 나면 다음 로직은 기존 스택 문제와 크게 다를 것이 없는데,

 

스택에 사람들을 넣어가며 새로운 사람이 기존 최상단 사람보다 키가 클 경우

 

볼 수 있는 사람의 쌍을 하나 증가시키고, 최상단 사람을 pop 하고

 

최상단 사람과 키가 같을 경우 최상단 사람의 count를 하나 올린 뒤 마찬가지로 pop 한다.

 

이렇게 반복하다가 스택 최상단 사람의 키가 새로운 사람의 키보다 크거나 스택이 비면 반복을 종료하고,

 

스택이 비어있지 않다면 그 사람은 무조건 한 명의 사람을 더 볼 수 있으니 볼 수 있는 사람의 쌍을 증가시킨 뒤

 

새로운 사람을 스택에 넣는다.

 

두 개의 필드를 가지는 객체를 스택에 넣는다는 발상이 굉장히 재미있었던 것 같다.

 

이어서 코드를 읽어보자.

		for (int i = 0; i < n; i++) {
			long height = Long.parseLong(br.readLine());
			long count = 1;

0부터 시작해서 반복문을 돈다.

 

먼저 새로운 사람의 키를 입력값으로 받고, count를 1로 초기화한다.

			while (!stack.isEmpty()) {
				if (stack.peek().height <= height) {
					result += stack.peek().count;
					if (stack.peek().height == height) {
						count += stack.peek().count;
					}
					stack.pop();

스택이 비어있지 않거나, 최상단 사람의 키가 새로운 사람의 키보다 작거나 같은 동안 반복한다.

 

만약 새로운 새로운 사람의 키가 최상단 사람의 키보다 클 경우

 

볼 수 있는 쌍을 저장하는 변수 result에 최상단 사람이 볼 수 있는 사람의 수를 더하는데,

 

이는 새로 들어오는 사람은 이 모든 사람을 볼 수 있기 때문이다.

 

이어서 만약 둘의 키가 같다면 새로운 사람의 count에 최상단 사람의 count를 더해준 뒤

 

최상단 사람을 제거(pop)한다.

				} else {
					break;
				}
			}

			if (!stack.isEmpty()) {
				result++;
			}

			stack.push(new Person(height, count));
		}

또한 탈출 조건이 성립해 반복문을 빠져나왔을 때 스택이 비어있지 않다면 역시 그 사람은 무조건

 

새로 들어오는 사람을 볼 수 있기 때문에 result에 1을 더해준다.

 

마지막으로 새로 들어온 사람을 스택에 넣고, 다음 인덱스로 넘어간다.

 

result를 출력하면 끝이다.

 

솔직히 이런 문제는 스스로 생각해서 풀 자신이 조금도 없다....

 

Java

 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayDeque;

public class Prob3015 {

	static class Person {
		long height;
		long count;

		public Person(long height, long count) {
			this.height = height;
			this.count = count;
		}
	}

	public static void main(String[] args) throws IOException {

		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

		int n = Integer.parseInt(br.readLine());
		ArrayDeque<Person> stack = new ArrayDeque<>();
		long result = 0;

		for (int i = 0; i < n; i++) {
			long height = Long.parseLong(br.readLine());
			long count = 1;

			while (!stack.isEmpty()) {
				if (stack.peek().height <= height) {
					result += stack.peek().count;
					if (stack.peek().height == height) {
						count += stack.peek().count;
					}
					stack.pop();
				} else {
					break;
				}
			}

			if (!stack.isEmpty()) {
				result++;
			}

			stack.push(new Person(height, count));
		}

		System.out.println(result);
	}
}

 

Python

 

import sys


class Person:
    def __init__(self, height, count):
        self.height = height
        self.count = count


n = int(sys.stdin.readline().rstrip())
stack = []
result = 0

for _ in range(n):
    height = int(sys.stdin.readline().rstrip())
    count = 1

    while stack:
        if stack[-1].height <= height:
            result += stack[-1].count
            if stack[-1].height == height:
                count += stack[-1].count
            stack.pop()
        else:
            break

    if stack:
        result += 1

    stack.append(Person(height, count))

print(result)

 

Performance