백준 3015번: 오아시스 재결합 (Stack)

문제 링크:3015번: 오아시스 재결합 (acmicpc.net)

3015번: 오아시스 재결합

이 문제는 만약 N명이 한 줄로 서있고 각 사람이 각각 의 키를 갖고 있을때, 서로를 볼 수 있는 쌍을 구하는 문제이다. 여기서 주어진 조건은

 

1. 키는 2^31보다 작은 숫자로 주어진다.

2. 두 사람 A,B 가 서로를 볼 수 있음을 판별할때는 만약 A,B 사이에  A키 또는 B키보다 큰 사람이 한명이라도 있으면, 볼수 없다.

 

처음 이 문제를 봤을때는 방향성을 못잡아서 미루어두었는데, 스위핑이나 스택 문제들을 풀어보고 어느정도 감을 잡을 수 있었다. 먼저 i (0<=i<N) 번째 사람이 들어간 쌍을 구할 떄, 쌍을 구하는 방법은 (j < i < k) 라고 했을때 i보다 작은 범위의 j를 이용해  (j, i)끼리 묶는 것과 큰 k를 이용해 (i, k) 로 묶는 경우 두개가 있을 것이다.

 

만약 모든 인덱스마다 앞 뒤로 모두를 조사한다면 중복되는 쌍이 있을 것이므로 나는 i번째 사람을 조사할 때에는 i보다 작은 사람들과 묶는 쌍 ((j, i) 의 경우) 을 조사하기로 하였다.

 

예시로 3, 5, 4, 1, 2 로 키가 주어지고 현재 index 3번째 (1) 을 조사한다고 했을때, index 0 ~ 2까지 최대값인 5보다 작은 값들은 어차피 가려지므로 조사할 필요가 없다. 그러므로 나는 순서대로 각 인덱스를 탐색하면서 임의의 stack에 내림차순이 되는 방법으로 교체하고 채워나가는 방법을 사용했다.

 

처음에는 스택이 비어있으므로

1번째는 [3]

2번째는 5가 스택의 마지막 값보다 크므로 pop해주고 [5]

3번째는 4가 스택의 마지막 값보다 작으므로 추가해서 [5,4]

4번쨰는 1가 스택의 마지막 값보다 작으므로 추가해서 [5,4,1]

5번째는 2가 스택의 마지막 값보다 크므로 pop 해주고, 마지막 값이 4가되므로 작으므로 추가해준다. [5,4,2]

 

이런 순서대로 스택을 저장해주었다. 여기서 pop해준다는 뜻은 자신보다 작으므로 무조건 볼 수 있다는 뜻이므로 pop해주는 횟수를 cnt에 올려 저장해주고, 마지막에 stack에 추가할 때 stack이 비어있지 않다면 마지막 값은 볼 수 있으므로 구한 cnt에 +1해주면 해당 인덱스에서 가능한 순서쌍의 개수를 구할 수 있다.

import sys

N=int(sys.stdin.readline().rstrip())

table=[0 for _ in range(N)]
stack=[]

def stackup(cur): #new가 stack[-1][0] 보다 작거나 같을때
    if cur<stack[-1][0]:
        stack.append([cur,0])
        return 1

    else: #new가 stack[-1][0]과 같을때
        last,same=stack[-1]
        
        if same==len(stack)-1: #stack 내 모든 사람이 중복일 때
            stack.append([cur,same+1])
            return same+1
        
        else: #stack 내 마지막 (same+1)명이 중복일때
            stack.append([cur,same+1])
            return same+2


def replace(cur): #new가 stack[-1][0] 보다 클 떄
    cnt=0
    while stack: #자신보다 작은사람 안나올때까지 pop, pop횟수저장
        if stack[-1][0]<cur:
            stack.pop()
            cnt+=1
            
        else: #중간에 크거나 같은사람 만난경우
            return cnt + stackup(cur)
        
    stack.append([cur,0])
    return cnt #stack내 모든 요소 pop 한경우
        
for i in range(N):
    new=int(sys.stdin.readline().rstrip())

    if not stack:
        stack.append([new,0]) #중복되는 키 사람들 처리 위해 [키,중복갯수-1] 형태로저장
        continue

    if new<=stack[-1][0]:
        table[i]=stackup(new)

    else:
        table[i]=replace(new)
        
print(sum(table))

기본적인 원리는 위와 같지만 여기서 같은 키를 가진 사람이 입력으로 주어질 수 있다. 같은 키를 가진 사람을 처리하기 위해서 stack에 [키, 중복인원] 형태로 저장해 나갔다. 만약 중복이 아니라면 [키,0] 1회 중복됐다면 (2명 같은키) [키,1] 이런 식으로 저장해 같은 키의 인원을 직전 값만 봐도 바로 알 수 있도록 하였다.

 

새로 추가되는 것이 stack의 마지막 키보다 작거나 같다면 바로 stackup() 해준다. stackup() 함수를 통해 만약 추가되는 값이 마지막 값보다 작다면 [new,0] 으로 바로 추가해주고 볼수 있는 사람은 직전의 한명이므로 return 1 해준다. 만약 마지막 값과 같다면 직전값의 same을 보아 모든 사람이 중복이라면 중복 시작 이전의 값이 없으므로 (same+1), 일부가 중복이라면 (same+1) + 1 (중복 시작 이전의 마지막 값) 을 return해준다.

 

replace()함수는 new 값이 stack의 마지막 키보다 클 때 시행해준다. stack을 pop해가며 stack의 마지막 키보다 new가 작거나 같아질 때까지 반복한다. pop하다가 만약 자신보다 큰 마지막 키를 발견했다면 여기서 기존의 stackup() 조건을 충족하므로 cnt  + strack() 을 리턴하여 반복을 줄였다. 만약 stack을 모두 pop했다면 cnt만큼 return해주면 된다. 

 

이렇게 하면 table에 각 인덱스를 조사할 때 마다 쌍의 개수가 저장된다. sum해주면 답이 나온다