[백준] 3015번 오아시스 재결합 - Java

문제 링크

https://www.acmicpc.net/problem/3015

문제

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

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

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

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

입력

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

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

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

출력

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

알고리즘: 자료구조, 스택

풀이

바로 이전에 풀었던 히스토그램 문제와 비슷한 유형이다.
먼저 조건을 보면 두 사람 사이에 A 또는 B보다 키가 큰 사람이 없어야 서로를 볼 수 있다.
그러면 A가 5 B가 6일 때 A 7 B는 서로 못 보지만, A 5 B는 볼 수 있다. 즉 A와 같은 키인 사람이 중간에 있어도 서로는 볼 수 있다. -> 이 조건 때문에 문제가 더 까다로워진다.

먼저 line에 같은 키가 없다고, 생각하고 문제에 접근해보자 일단 N의 범위는 최대가 500,000이기 때문에 O(N)풀이를 생각해야 한다.
사람들이 내림차순으로 서 있다고 생각해보자, 5 4 3 2 1 -> 첫 번째 사람을 제외하고는 모두 앞 사람만을 볼 수 있다. 이 규칙을 찾아냈으면 stack을 이용해서 stack이 내림차순을 유지하면서 쌍을 계산하면 된다. 여기서 stack에 5 4 3 2 1이 존재하는데 stack에 top보다 큰 값이 들어오려고 한다면, stack이 내림차순을 유지하게끔 현재 키보다 작은 값들을 pop 해주면 된다.
pop을 할 때 쌍을 계산해주는 게 포인트다. 어떻게 쌍을 계산하냐? 먼저 현재 키와 pop해주는 키는 한 쌍이기 때문에 +1 해준다. 그리고 pop해주는 키 앞에 값이 존재한다면 그 쌍 또한 +1 해주면 끝난다. stack안에는 내림차순을 유지하고 있기 때문에 인접한 사람 외에는 서로를 볼 수 없다.

여기서 키가 같은 사람이 존재하면 이야기가 달라진다. 키가 같다면 5 2 2 맨 마지막에 2는 자기 앞에 2와 그 앞인 5를 볼 수 있다.
같은 키가 존재하는 line은 어떻게 풀어야 하냐?
해답은 같은 키가 몇 명인지 알 수 있는 정보도 같이 stack에 넣어주면(5 2 2는 -> 5는 1개, 2는 2개) 쉽게 풀린다.
똑같이 stack이 내림차순을 유지하게끔 넣어주고, 마냐게 stack에 top과 현재 키가 같다면 해당 키에 cout + 1 해주고, 현재 키가 top보다 큰 값이 들어왔다면 stack을 내림차순으로 유지해야 되기 때문에 현재 키보다 작은 값을 모두 pop 해준다. pop 해줄 때 쌍을 계산하는데 해당 키의 개수는 쌍의 개수이기 때문에 그 만큼 + 해주고 (ex (5 2 2) 6 -> 6 2, 6 2) check 해주는 키 앞에 값이 존재한다면 키의 cout를 -1 해주면서 더해주면 되고, 키 앞에 값이 존재하지 않다면 키의 cout를 -1 해주고 그 값에 -1 해주면서 더해주면 된다. -> 0이 될 때까지
예를 들면 stack -> (5, 2, 2, 2) 현재 키 4라고 할 때 4 2, 4 2, 4 2 이렇게 2를 가진 사람의 수만큼 더해주고, (5,2(3)) 2앞에 5가 존재하기 때문에 맨 뒤에 존재하는 2는 (2 2, 2 2, 2 5) 3개의 쌍을 그 다음 2는 (2 2, 2 5) 2개의 쌍을 마지막 2는 (2 5) 1개의 쌍을 가지고 있으므로 3 + 2 + 1 쌍을 가진다. 만약 2앞에 5가 존재하지 않는다면 (2 2 2) 맨 뒤에 2는 (2 2, 2 2) 2개의 쌍을 그다음 2는 (2 2) 1개의 쌍으로 총 2 + 1개의 쌍을 가진다.
그래서 요약하면 2의 개수가 3개일 때 앞에 값이 존재하면 3 + 2 + 1, 값이 존재하지 않는다면 2 + 1이라는 규칙을 가지고 쌍을 계산해주면 된다.

소스 코드

import java.io.*;
import java.util.*;

class Node {
    int h,c;
    Node(int h, int c) {
        this.h = h;
        this.c = c;
    }
}

public class Main {
    static int N;
    static Stack<Node> stack = new Stack<>();
    static long ans = 0;
    public static void main(String args[]) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        N = Integer.parseInt(br.readLine());
        for(int i=0; i<N; i++) {
            int height = Integer.parseInt(br.readLine());
            if(stack.empty()) stack.push(new Node(height, 1));
            else {
                if(stack.peek().h > height) stack.push(new Node(height, 1));
                else if(stack.peek().h == height) stack.peek().c += 1;
                else {
                    while(!stack.empty()) {
                        if(stack.peek().h > height) {
                            stack.push(new Node(height, 1));
                            break;
                        } else if(stack.peek().h == height) {
                            stack.peek().c += 1;
                            break;
                        } else {
                            Node n = stack.pop();
                            ans += n.c;
                            if(stack.empty()) n.c -= 1;
                            for(int j=n.c; j>0; j--) ans += j;
                        }
                    }
                    if(stack.empty()) stack.push(new Node(height, 1));
                }
            }
        }
        //stack이 비어 있지 않다면
        while(!stack.empty()) {
            Node n = stack.pop();
            if(stack.empty()) n.c -= 1;
            for(int i=n.c; i>0; i--) ans += i;
        }
        System.out.println(ans);
    }
}