파이썬 코테 알아야할 것

들어가며

파이썬은 코테의 적폐언어 중 하나이다. 그만큼 내부적으로 지원하는 함수가 많기 때문에 코딩테스트 진행시 다른 언어 보다 유리하다. 그렇기 때문에 파이썬을 제대로 활용하기 위해서 알아야할 몇가지를 정리한다.

itertools의 조합형 함수

combinations(iterable, r) 함수는 주어진 iterable에서 원소 r개를 뽑아 만들 수 있는 조합을 생성합니다. 순서는 고려하지 않습니다.

# 예시: [1, 2, 3]에서 2개를 뽑는 조합
combs = list(combinations([1, 2, 3], 2))
print(combs)
# 출력: [(1, 2), (1, 3), (2, 3)]

permutations(iterable, r) 함수는 주어진 iterable에서 원소 r개를 뽑아 만들 수 있는 순열을 생성합니다. 순서를 고려합니다.

# 예시: [1, 2, 3]에서 2개를 뽑는 순열
perms = list(itertools.permutations([1, 2, 3], 2))
print(perms)
# 출력: [(1, 2), (1, 3), (2, 1), (2, 3), (3, 1), (3, 2)]

product(*iterables, repeat=1) 함수는 여러 iterable의 데카르트 곱(모든 가능한 조합)을 생성합니다.

# 예시: [1, 2]와 ['A', 'B']의 데카르트 곱
prod = list(itertools.product([1, 2], ['A', 'B']))
print(prod)
# 출력: [(1, 'A'), (1, 'B'), (2, 'A'), (2, 'B')]

# 동일 iterable의 반복된 데카르트 곱
repeat_prod = list(itertools.product([1, 2], repeat=2))
print(repeat_prod)
# 출력: [(1, 1), (1, 2), (2, 1), (2, 2)]

combinations_with_replacement(iterable, r) 함수는 주어진 iterable에서 원소 r개를 중복 허용하여 뽑아 만들 수 있는 조합을 생성합니다.

# 예시: [1, 2, 3]에서 2개를 중복 허용하여 뽑는 조합
combs_wr = list(itertools.combinations_with_replacement([1, 2, 3], 2))
print(combs_wr)
# 출력: [(1, 1), (1, 2), (1, 3), (2, 2), (2, 3), (3, 3)]

set은 중복되지 않는 원소들의 집합을 나타내며, 원소의 존재 여부를 빠르게 확인할 수 있습니다.

set은 다음과 같은 문제 유형에서 자주 사용됩니다:
중복 제거: 주어진 리스트에서 중복된 원소를 제거해야 할 때.
원소 존재 여부 확인: 특정 원소가 리스트나 다른 자료구조에 존재하는지 빠르게 확인해야 할 때.
교집합, 합집합, 차집합: 두 개 이상의 집합에 대해 교집합, 합집합, 차집합을 구할 때.

# 중복 제거 예시
def remove_duplicates(lst):
    return list(set(lst))

# 공통 원소 찾기 예시
def find_common_elements(lst1, lst2):
    return list(set(lst1) & set(lst2))

# set 생성 - O(1)
my_set = {1, 2, 3, 4}
print(my_set)  # 출력: {1, 2, 3, 4}

# 원소 추가 - O(1) 평균
my_set.add(5)
print(my_set)  # 출력: {1, 2, 3, 4, 5}

# 원소 존재 여부 확인 - O(1) 평균
print(3 in my_set)  # 출력: True

# 원소 제거 - O(1) 평균
my_set.remove(2)
print(my_set)  # 출력: {1, 3, 4, 5}

bisect 모듈은 정렬된 리스트에 대해 이진 탐색 및 삽입 작업을 지원합니다. bisect는 기본적으로 이진 탐색을 수행하여 원소가 삽입될 위치를 찾습니다.

bisect은 다음과 같은 문제 유형에서 자주 사용됩니다:
이진 탐색: 정렬된 리스트에서 특정 원소의 삽입 위치나 존재 여부를 확인할 때.
효율적인 삽입: 정렬된 리스트에 원소를 효율적으로 삽입할 때.

import bisect

# 정렬된 리스트에 삽입 예시
def insert_sorted(lst, x):
    bisect.insort(lst, x)
    return lst

# 특정 값보다 큰 첫 번째 원소 찾기 예시
def find_first_greater(lst, x):
    idx = bisect.bisect_right(lst, x)
    if idx < len(lst):
        return lst[idx]
    return None

# 정렬된 리스트
sorted_list = [1, 2, 4, 4, 5]

# 이진 탐색하여 삽입할 위치 찾기 - O(log n)
index = bisect.bisect(sorted_list, 3)
print(index)  # 출력: 2

# 원소 삽입 - O(n)
bisect.insort(sorted_list, 3)
print(sorted_list)  # 출력: [1, 2, 3, 4, 4, 5]

heapq 모듈은 힙(우선순위 큐) 알고리즘을 제공하며, 최소 힙을 기본으로 합니다. 최대 힙을 구현하려면 원소의 부호를 반대로 하여 삽입/제거할 수 있습니다.

heapq은 다음과 같은 문제 유형에서 자주 사용됩니다:
우선순위 큐: 작업의 우선순위를 정하고, 높은 우선순위부터 처리해야 할 때.
최소값/최대값 찾기: 리스트에서 최소값 또는 최대값을 빠르게 찾아야 할 때

import heapq

# 리스트를 힙으로 변환 - O(n)

# 우선순위 큐 예시
def process_tasks(tasks):
    heapq.heapify(tasks)
    result = []
    while tasks:
        result.append(heapq.heappop(tasks))
    return result

# 상위 k개의 최소값 찾기 예시
def find_k_smallest(lst, k):
    return heapq.nsmallest(k, lst)

heap = [4, 1, 7, 3, 8, 5]
heapq.heapify(heap)
print(heap)  # 출력: [1, 3, 5, 4, 8, 7]

# 원소 추가 - O(log n)
heapq.heappush(heap, 2)
print(heap)  # 출력: [1, 2, 5, 3, 8, 7, 4]

# 최소값 제거 - O(log n)
min_element = heapq.heappop(heap)
print(min_element)  # 출력: 1
print(heap)  # 출력: [2, 3, 5, 4, 8, 7]

deque는 양방향 큐로, 양쪽에서 빠르게 삽입과 삭제를 할 수 있습니다.

deque은 다음과 같은 문제 유형에서 자주 사용됩니다:
양방향 큐: 큐의 양쪽 끝에서 삽입과 삭제가 빈번하게 일어날 때.
슬라이딩 윈도우: 고정 크기의 윈도우를 사용하여 리스트를 순회할 때.

from collections import deque

# 슬라이딩 윈도우 예시
def sliding_window_sum(lst, k):
    dq = deque(lst[:k])
    result = [sum(dq)]
    for i in range(k, len(lst)):
        dq.popleft()
        dq.append(lst[i])
        result.append(sum(dq))
    return result

# 양방향 큐 예시
def deque_operations():
    dq = deque()
    dq.append(1)
    dq.appendleft(2)
    dq.pop()
    dq.popleft()
    return dq

# deque 생성 - O(n)
dq = deque([1, 2, 3, 4])
print(dq)  # 출력: deque([1, 2, 3, 4])

# 오른쪽에 원소 추가 - O(1)
dq.append(5)
print(dq)  # 출력: deque([1, 2, 3, 4, 5])

# 왼쪽에 원소 추가 - O(1)
dq.appendleft(0)
print(dq)  # 출력: deque([0, 1, 2, 3, 4, 5])

# 오른쪽에서 원소 제거 - O(1)
dq.pop()
print(dq)  # 출력: deque([0, 1, 2, 3, 4])

# 왼쪽에서 원소 제거 - O(1)
dq.popleft()
print(dq)  # 출력: deque([1, 2, 3, 4])

Counter는 요소의 개수를 셀 때 유용한 dict의 서브클래스입니다.

Counter는 다음과 같은 문제 유형에서 자주 사용됩니다:
빈도 계산: 리스트나 문자열에서 각 원소의 빈도를 계산해야 할 때.
가장 빈번한 요소 찾기: 가장 자주 등장하는 원소를 찾아야 할 때.

from collections import Counter

# 가장 빈번한 문자 찾기 예시
def most_common_char(s):
    counter = Counter(s)
    return counter.most_common(1)[0]

# 가장 빈번한 k개의 요소 찾기 예시
def most_common_elements(lst, k):
    counter = Counter(lst)
    return counter.most_common(k)

# 리스트 요소의 개수를 셈 - O(n)
count = Counter(['apple', 'banana', 'apple', 'orange', 'banana', 'apple'])
print(count)  # 출력: Counter({'apple': 3, 'banana': 2, 'orange': 1})

# 특정 요소의 개수 확인 - O(1)
print(count['apple'])  # 출력: 3

# 요소 추가 - O(k) (k는 추가되는 요소의 개수)
count.update(['banana', 'orange'])
print(count)  # 출력: Counter({'apple': 3, 'banana': 3, 'orange': 2})