코딩 테스트에 유용한 파이썬 라이브러리와 함수들

개인적으로 노션에 정리해두고 생각이 나지 않을 때 참고하던 내용을 블로그에 공유하고자 한다.

코딩 테스트에 따라서는 사용할 수 없는 라이브러리가 존재할 수 있으니 반드시 확인해보아야 한다.

---

itertools

permutations

n 개의 데이터를 뽑아 나열 (순열)

from itertools import permutations

data = ['A', 'B', 'C']
result = list(permutations(data, 3))

# [("A", "B", "C"), ("A", "C", "B"), ("B", "A", "C"), ("B", "C", "A"), ("C", "A", "B"), ("C", "B", "A")]

 

combinations

n 개의 데이터를 뽑아 순서 없이 나열 (조합)

from itertools import combinations

data = ['A', 'B', 'C']
result = list(combinations(data, 3))

# [('A', 'B'), ('A', 'C'), ('B', 'C)]

 

product

n 개의 데이터를 뽑아 일렬로 나열하는 모든 경우 (순열)

단, 원소를 중복해서 뽑는다.

from itertools import product

data = ['A', 'B', 'C'] # 데이터 준비
result = list(product(data, repeat=2)) # 2개를 뽑는 모든 순열 구하기(중복 허용)

# [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'A'), ('B', 'B'), ('B', 'C'), ('C', 'A'), ('C', 'B'), ('C', 'C')]

 

combinations_with_replacement

n 개의 데이터를 뽑아 순서를 고려하지 않고 나열 (조합)

단, 원소를 중복해서 뽑는다.

from itertools import combinations_with_replacement

data = ['A', 'B', 'C'] # 데이터 준비
result = list(combinations_with_replacement(data, 2)) # 2개를 뽑는 모든 조합 구하기(중복허용)

# [('A', 'A'), ('A', 'B'), ('A', 'C'), ('B', 'B'), ('B', 'C'), ('C', 'C')]

 

compress

주어진 iterable 객체의 원소 중 selectors에서 같은 인덱스의 원소가 참인 것만 뽑아 만든 iterator 리턴

from itertools import compress

list(compress('ABCDEF', [1, 0, 1, 0, 1, 1]))
# 두 번째 매개변수가 selectors 이다.
# ['A', 'C', 'E', 'F']

 

---

 

heapq

힙을 이용한 우선순위 큐 기능 구현에 사용

PriorityQueue 보다 빠름

힙을 사용하면 O(NlogN)의 복잡도로 오름차순 정렬 가능

파이썬에서는 최대 힙을 제공하지 않음

• heapq.heapify(리스트)
  • 리스트를 힙 구조로 변경
• heapq.heappush(리스트, 요소)
  • 리스트에 힙 구조를 유지하며 요소 삽입
• heapq.heappop(리스트)
  • 리스트에 힙 구조를 유지하며 요소 꺼내기

import heapq

def heapsort(iterable):
	h = []
	result = []
	
	# 모든 원소 힙에 삽입
	for value in iterable:
		heapq.heappush(h, value)

	# 힙에 삽입된 모든 원소 차례로 꺼내서 담기
	for i in range(len(h)):
		result.append(heapq.heappop(h))
	
	return result

result = heapsort([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])
print(result)

 

TIP: 최대 힙으로 사용하고 싶으면 원소를 삽입할 때 -1을 곱하고, 원소를 꺼낸 후 다시 -1을 곱해서 원상 복구한다. 단, 숫자가 아닌 타입일 때는 사용할 수 없는 방법이다.

TIP2: (a, b, c)와 같이 리스트나 튜플을 원소로 사용하면 a, b, c 순으로 비교한다. a가 같다면 b와, b와 같다면 c와 비교한다.

 

---

 

bisect

이진 탐색을 위한 라이브러리

정렬된 배열에서 특정 원소를 찾을 때 유리하다.

O(logN)의 탐색 속도

 

• bisect_left(a, x)
  • 정렬된 순서를 유지하면서 리스트 a에서 데이터x를 삽입할 가장 왼쪽 인덱스를 찾는 메서드
• bisect_right(a, x)
  • 정렬된 순서를 유지하면서 리스트 a에서 데이터 x를  삽입할 가장 오른쪽 인덱스를 찾는 메서드

from bisect import bisect_left, bisect_right

a = [1, 2, 4, 4, 8]
x = 4

print(bisect_left(a, x)) # 2
print(bisect_right(a, x)) # 4

 

---

Collections

deque

리스트에서 첫 번째 요소를 제거하거나 첫 번째 위치에 요소를 추가하는 경우 O(N) 이 소요되는 것에 반해 deque은 O(1) 만큼이 소요된다. 이러한 특징으로 Queue 구조가 필요할 때 유용하다.

단, 리스트와 다르게 인덱싱, 슬라이싱 등은 지원되지 않는다. (LinkedList 구조를 사용하기 때문)

from collections import deque

data = deque([2, 3, 4])
data.appenleft(1)
data.append(5)

 

counter

등장 횟수를 세는 기능

iterable 객체가 주어질 때 해당 객체 내부 원소가 몇 번 등장했는지 알 수 있다.

from collections import Counter

counter = Counter(['red','blue','red','green','blue','blue'])

print(counter['blue']) # 3
print(counter['green']) # 1
print(dict(counter)) # 딕셔너리로 변환
# {'red' : 2, 'blue': 3, 'green': 1}

 

math

• math.factorial(n)
  • n! (팩토리얼)
• math.sqrt(n)
  • n의 제곱근
• math.gcd([a, b, c, ...])
  • 최대 공약수
• math.pi
• math.e

 

---

유용한 파이썬 함수

divmod

몫과 나머지를 한 번에 처리

매우 큰 수에서 연산자보다 좀 더 빠름

q, r = divmod(10, 3)

 

zip

주어지는 iterable 객체들의 각 요소를 묶어 tuple의 iterator로 반환

list(zip('ABCD', 'xy'))
# 원소의 개수가 적은 'xy'에 맞춰 tuple이 2개만 생성됨
# [('Ax', 'By')]

 

set

iterable 객체를 set으로 변환해준다. (set은 중복 없이 저장하므로 중복을 제거하는데 유용하게 사용된다)

set은 hashable한 값만 받을 수 있다. (리스트 불가능, 튜플 가능)

set([1, 2, 3])
# [1, 2, 3]의 각 원소가 int 타입이므로 set에 넣을 수 있음

set([[1, 2, 3]])
# [[1, 2, 3]]의 원소인 [1, 2, 3]은 list 타입 변수이므로 에러!

 

• 합집합
  • a.union(b)
  • a | b
• 교집합
  • a.intersection(b)
  • a & b
• 차집합
  • a.difference(b)
  • a - b

단, a와 b는 Set

 

• issubset
  • 부분 집합이면 True 반환
• issuperset
  • 상위 집합이면 True 반환
• isdisjoint
  • 두 집합 간 교집합이 없으면 True 반환

set([1, 2]).issubset(set([1, 2, 3]))
# True
set([1, 2, 3]).issuperset(set([1, 2]))
# True
set([1, 2, 3]).isdisjoint(set([4, 5]))
# True