collections, heapq를 사용해보자

파이썬을 이용한 코딩테스트 연습시 원하고자 하는 기능을 리스트와 딕셔너리로 구현할때 시간초과나 메모리초과라는 문제점을 만날때가 있다.

error_image

1. 카운팅 기능을 리스트에서 구현할때

count 기능을 for문 사용시

#count
words = [0,1,2,3,4,2,3,1,2,4,0]
for i in range(len(words)):
    print(words.count(i),end=" ")
# 출력결과 > 1 2 3 2

이때의 시간복잡도는 O(N)인데 for문 안에서 n번 반복하기 때문에 O(N^2)이 됩니다.

하지만 collections의 Counter를 사용시

#Counter
from collections import Counter
words = [0,1,2,3,4,2,3,1,2,4,0]
counter = Counter(words)
for i in range(len(words)):
    print(counter[i],end=" ")

시간복잡도는 O(1) for문을 통해 n번 반복하기 때문에 O(N)이 된다.

또한 딕셔너리와 같은 기능을 제공하기 때문에 앞으로 카운팅 기능이 필요할때 딕셔너리에 저장하는 것보다 그냥 collections의 Counter기능을 쓸때 시간복잡도를 더 줄일 수 있다.

2. 리스트의 공간복잡도가 커진다면?

queue 기능을 파이썬에서 구현하고자 한다면 heapq를 사용하면 좋다 예를 들어 리스트에서 sort시켜 맨 앞의 작은 수끼리 더해나가는 자료구조를 생각할때 리스트로 구현하면 메모리초과가 날 수 있다.

# 1715, 카드 정렬하기
import sys

n = int(sys.stdin.readline())
ns = list(map(int,sys.stdin.read().split()))
ns.sort()
temp = 0
for i in range(len(ns)-1):
    temp += ns[i] + ns[i+1]
    ns[i+1] = temp
print(temp)

# 1715, 카드 정렬하기
import sys
import heapq


n = int(sys.stdin.readline())
ns = list(map(int,sys.stdin.read().split()))
if len(ns) == 1:
    print(0)
else:
    ans = 0
    while len(ns) > 1:
        temp1 = heapq.heappop(ns)
        temp2 = heapq.heappop(ns)
        ans += temp1 + temp2
        heapq.heappush(ns, temp1 + temp2)
    print(ans)

리스트를 통해 구현할때보다 메모리를 아낄 수 있다.