[백준] 17299 - 오등큰수 in 파이썬

문제 출처

17299번: 오등큰수

 

문제

크기가 N인 수열 A = A1, A2, ..., AN이 있다. 수열의 각 원소 Ai에 대해서 오등큰수 NGF(i)를 구하려고 한다.

Ai가 수열 A에서 등장한 횟수를 F(Ai)라고 했을 때, Ai의 오등큰수는 오른쪽에 있으면서 수열 A에서 등장한 횟수가 F(Ai)보다 큰 수 중에서 가장 왼쪽에 있는 수를 의미한다. 그러한 수가 없는 경우에 오등큰수는 -1이다.

예를 들어, A = [1, 1, 2, 3, 4, 2, 1]인 경우 F(1) = 3, F(2) = 2, F(3) = 1, F(4) = 1이다. A1의 오른쪽에 있으면서 등장한 횟수가 3보다 큰 수는 없기 때문에, NGF(1) = -1이다. A3의 경우에는 A7이 오른쪽에 있으면서 F(A3=2) < F(A7=1) 이기 때문에, NGF(3) = 1이다. NGF(4) = 2, NGF(5) = 2, NGF(6) = 1 이다.

 

입력

첫째 줄에 수열 A의 크기 N (1 ≤ N ≤ 1,000,000)이 주어진다. 둘째에 수열 A의 원소 A1, A2, ..., AN (1 ≤ Ai ≤ 1,000,000)이 주어진다.

출력

총 N개의 수 NGF(1), NGF(2), ..., NGF(N)을 공백으로 구분해 출력한다.

---

문제 풀이

오등큰수 NGF(i) 란 다음과 같다.

1. A[i] 의 오른쪽에 있으면서 (i+1, i+2,...)
2. 등장 횟수 또한 더 커야 한다.

 

 즉, 수열 A의 원소 각각의 등장횟수를 구하고, A[i]의 등장횟수와 A[i]의 오른쪽에 있는 원소의 등장횟수를 비교하면 된다.

 오등큰수를 구하지 못했을 경우에는 해당 인덱스를 스택에 push하고, 오등큰수 조건이 성립한다면 인덱스가 저장된 스택을 pop하면서 답을 기록한다.

 0인덱스의 오등큰수를 구하지 못하고, 1인덱스의 오등큰수도 구하지 못하다가, 2인덱스에서 1인덱스의 오등큰수를 발견했다면, 결과적으로 0인덱스와 1인덱스는 같은 오등큰수를 가지기 때문에 스택의 성질을 이용한다.

 

1. 수열의 크기 N을 입력받는다.
2. 수열 A 를 리스트로 입력받는다.
3. 1,000,001 크기의 F 배열에 수열 A의 원소 각각의 등장횟수를 구한 뒤 저장한다. (수열 A의 원소 : 1 ≤ Ai ≤ 1,000,000)
4. 답을 기록할 배열을 -1로 초기화한다. (오등큰수를 찾지 못했을 경우 -1 로 표현하기로 약속했기 때문)
5. 스택을 초기화한다. (스택에는 인덱스 번호를 넣는다)
   • 수열 A의 첫번째 인덱스부터 오등큰수를 찾아야 하므로, 첫번째 인덱스인 0을 push한다.
6. 수열 A를 순회하면서 오등큰수를 구한다.
   1. 오등큰수는 오른쪽에 있는 원소에 있으므로 순회는 두번째 인덱스부터 시작한다.
   2. 스택의 top 인덱스의 값 등장횟수보다 순회중인 인덱스의 값이 등장횟수가 더 크다면, 오등큰수를 구한 것이므로 스택을 pop한다.
   3. 오등큰수를 찾는 연산이 끝났거나, 오등큰수를 해당 순회에서 찾지 못했다면, 해당 인덱스를 스택에 push하여 다음 순회를 준비한다.

 

 

작성 코드

import sys, copy
input = sys.stdin.readline

N = int(input())
A = list(map(int,input().split()))
F = [0] * 1000001
st = copy.deepcopy(A)
ans = [-1] * N

while st:
    F[st.pop()] += 1

st.append(0)
for i in range(1, N):
    while st and F[A[st[-1]]] < F[A[i]]:
        ans[st.pop()] = A[i]
    st.append(i)

print(*ans) # python3: 2088ms / pypy3: 740ms

제출 결과 python3는 2088ms, pypy3는 740ms 로, pypy3가 더 빠른 결과를 보였다.

 

 

등장횟수를 세는 F 배열의 인덱스 주의!!!

 처음에 인덱스 에러가 나서 다시 찬찬히 살펴보니, 어이없는 실수를 했다. 1000000 크기로 선언해줬던 것.

수열 A의 원소는 1 ≤ Ai ≤ 1,000,000 범위를 갖는다. F배열은 Ai의 등장횟수를 저장하므로, 0번째 인덱스는 버리는 인덱스가 된다. 따라서 배열의 크기는 1000000+1 인 1000001 로 선언해줘야한다.

 

 

---

* 다른사람들은 어떻게 작성했는지 탐색해보던 중 알게된 것

직접 등장횟수를 세지 않아도 Counter 라는 파이썬 컬렉션을 사용하면 간단하게 구할 수 있다는 것을 알았다.

Counter 컬렉션의 쓰임을 알아보고 정리해보는 시간을 가져야겠다.