[Silver IV] 배열 전체 탐색하기 - 26168

문제 링크

성능 요약

메모리: 33320 KB, 시간: 136 ms

분류

이분 탐색, 정렬

문제 설명

크기 n인 정수형 배열 A가 주어진다. 배열 A의 원소는 A[0], A[1], ... , A[n-1]이다. 배열 A에는 같은 값을 갖는 원소가 여러 개 존재할 수 있다. 배열 A에 대한 m개의 질의가 저장된 배열 B가 주어진다. 배열 B에 저장된 m개의 질의는 아래 세 가지 유형으로 구분된다. 첫 번째가 유형 1, 두 번째가 유형 2, 세 번째가 유형 3이다.

• 1 k: 배열 A의 원소 중 k보다 크거나 같은 원소의 개수를 출력한다.
• 2 k: 배열 A의 원소 중 k보다 큰 원소의 개수를 출력한다.
• 3 i j: 배열 A의 원소 중 i보다 크거나 같고 j보다 작거나 같은 원소의 개수를 출력한다.

배열 B에 저장된 첫 번째 질의부터 m번째 질의까지 순서대로 처리하면서 질의 결과를 출력하자.

입력

첫 번째 줄에 n과 m이 공백을 사이에 두고 순서대로 주어진다.

두 번째 줄에 배열 A의 원소 A[0], A[1], ... , A[n-1]이 공백을 사이에 두고 순서대로 주어진다.

세 번째 줄부터 m개의 줄에 걸쳐 배열 B에 저장된 m개의 질의가 순서대로 주어진다. 한 줄에 하나의 질의를 나타내는 정수가 공백을 사이에 두고 순서대로 주어진다.

출력

첫 번째 줄부터 질의 결과를 순서대로 한 줄씩 출력한다.

---

풀이

아이디어

bisect 모듈을 이용해 이분탐색을 구현하기로 했다.
아래는 bisect에서 사용하는 함수들이다.

함수	설명
bisect.bisect_left(target, x)	target 리스트에서 x가 들어갈 가장 왼쪽의 인덱스를 탐색
bisect.bisect_right(target, x)	target 리스트에서 x가 들어갈 가장 오른쪽의 인덱스를 탐색

위 두 특성을 이용하여 약간의 응용을 거쳤다.

문제에서 요구하는 기능은 아래 3가지이다.

1. k 보다 크거나 같은 원소의 개수.

len(li) - bisect.bisect_right(li, x-1)

2. k 보다 큰 원소의 개수.

len(li) - bisect.bisect_right(li, x)

3. i와 j 사이의 원소의 개수.

bisect.bisect_right(li, y) - bisect.bisect_left(li, x)

코드

import bisect
import sys
input = sys.stdin.readline
print = sys.stdout.write
def more_then(li, x):
    return len(li) - bisect.bisect_right(li, x-1)

def more(li, x):
    return len(li) - bisect.bisect_right(li, x)

def between(li, x, y):
    return bisect.bisect_right(li, y) - bisect.bisect_left(li, x)
n, m = map(int, input().split())
lst = sorted(list(map(int, input().split())))
for i in range(m):
    option, *t = map(int, input().split())
    if option == 1:
        print(str(more_then(lst, *t)) + "\n")
    
    elif option == 2:
        print(str(more(lst, *t)) + "\n")
    
    else:
        print(str(between(lst, *t)) + "\n")

---

회고

이분탐색의 개념을 알았으니 시간 단축을 위해 내장 모듈인 bisect를 이용해 풀이한 첫번째 문제이다. 조금 어색했지만 시간은 확실히 단축할 수 있었다.

처음에는 시간초과가 나왔는데, 깜짝 놀랐었다. 단순히 입력 출력이 많아 거기서 시간초과가 난 것이므로 sys 모듈을 이용하여 해결하였다.