정렬과 탐색 코드
노트의 개념을 Python 3로 옮긴 코드다. 표준 라이브러리만 쓰고, 파일 끝의 assert 예제가 그대로 실행된다. 복사해 python3 파일이름.py로 돌려 보자.
버블 정렬
"""버블 정렬: 이웃한 두 칸을 비교해 순서가 틀리면 바꾼다.
시간: 최선 O(n) (이미 정렬된 입력, 조기 종료), 평균·최악 O(n²)
공간: O(1) — 배열 안에서 바꾼다. 안정 정렬이다.
"""
def bubble_sort(values: list) -> int:
"""values를 제자리에서 오름차순으로 정렬하고, 두 칸을 바꾼 횟수를 돌려준다."""
swaps = 0
for end in range(len(values) - 1, 0, -1):
swapped_in_pass = False
for index in range(end):
if values[index] > values[index + 1]:
values[index], values[index + 1] = values[index + 1], values[index]
swaps += 1
swapped_in_pass = True
# 한 바퀴 동안 바꾼 일이 없으면 이미 정렬된 것이다
if not swapped_in_pass:
break
return swaps
def count_inversions(values: list) -> int:
"""앞이 더 큰 쌍(역순 쌍)의 수. 버블 정렬의 바꾸기 횟수와 같다."""
return sum(
1
for left in range(len(values))
for right in range(left + 1, len(values))
if values[left] > values[right]
)
if __name__ == "__main__":
data = [4, 2, 5, 1, 3]
inversions = count_inversions(data)
assert bubble_sort(data) == inversions == 6
assert data == [1, 2, 3, 4, 5]
already_sorted = [1, 2, 3]
assert bubble_sort(already_sorted) == 0
assert bubble_sort([]) == 0
print("ok")선택·삽입 정렬
"""선택 정렬과 삽입 정렬: 왼쪽에 정렬된 묶음을 한 칸씩 늘린다.
시간: 선택 정렬은 늘 O(n²) 비교에 바꾸기 많아야 n-1번,
삽입 정렬은 최선 O(n) (이미 정렬된 입력), 평균·최악 O(n²)
공간: O(1) — 배열 안에서 옮긴다. 삽입 정렬은 안정, 선택 정렬은 불안정하다.
"""
import random
def selection_sort(values: list) -> int:
"""values를 제자리에서 오름차순으로 정렬하고, 두 칸을 바꾼 횟수를 돌려준다."""
swaps = 0
for start in range(len(values) - 1):
min_index = start
for index in range(start + 1, len(values)):
if values[index] < values[min_index]:
min_index = index
if min_index != start:
values[start], values[min_index] = values[min_index], values[start]
swaps += 1
return swaps
def insertion_sort(values: list) -> int:
"""values를 제자리에서 오름차순으로 정렬하고, 큰 값을 오른쪽으로 민 횟수를 돌려준다."""
shifts = 0
for next_index in range(1, len(values)):
key = values[next_index]
hole = next_index
# 같은 값은 넘지 않아야(>만 민다) 안정 정렬이 된다
while hole > 0 and values[hole - 1] > key:
values[hole] = values[hole - 1]
hole -= 1
shifts += 1
values[hole] = key
return shifts
if __name__ == "__main__":
data = [5, 2, 4, 6, 1, 3]
assert insertion_sort(data) == 9
assert data == [1, 2, 3, 4, 5, 6]
reversed_data = [6, 5, 4, 3, 2, 1]
assert selection_sort(reversed_data) == 3
assert reversed_data == [1, 2, 3, 4, 5, 6]
assert insertion_sort([6, 5, 4, 3, 2, 1]) == 15
assert insertion_sort([1, 2, 3]) == 0
assert selection_sort([]) == 0
assert insertion_sort([7]) == 0
rng = random.Random(0)
for _ in range(200):
sample = [rng.randint(0, 9) for _ in range(rng.randint(0, 12))]
for sort in (selection_sort, insertion_sort):
copy = list(sample)
sort(copy)
assert copy == sorted(sample)
print("ok")병합 정렬
"""병합 정렬: 반으로 나눠 각각 정렬한 뒤, 정렬된 두 조각을 앞에서부터 비교하며 합친다.
시간: 최선·평균·최악 모두 O(n log n)
공간: O(n) — 합칠 때 결과를 담을 임시 리스트가 필요하다. 안정 정렬이다.
"""
import random
def merge(left: list, right: list) -> tuple[list, int]:
"""정렬된 두 리스트를 합친 리스트와 비교 횟수를 돌려준다. 같으면 left를 먼저 꺼내 안정성을 지킨다."""
merged = []
comparisons = 0
left_index = right_index = 0
while left_index < len(left) and right_index < len(right):
comparisons += 1
if left[left_index] <= right[right_index]:
merged.append(left[left_index])
left_index += 1
else:
merged.append(right[right_index])
right_index += 1
# 한쪽이 바닥나면 남은 쪽은 이미 정렬돼 있어 비교 없이 붙인다
merged.extend(left[left_index:])
merged.extend(right[right_index:])
return merged, comparisons
def merge_sort(values: list) -> tuple[list, int]:
"""values를 정렬한 새 리스트와 전체 비교 횟수를 돌려준다."""
if len(values) <= 1:
return list(values), 0
middle = len(values) // 2
left, left_comparisons = merge_sort(values[:middle])
right, right_comparisons = merge_sort(values[middle:])
merged, merge_comparisons = merge(left, right)
return merged, left_comparisons + right_comparisons + merge_comparisons
if __name__ == "__main__":
assert merge([1, 4, 7, 8], [2, 3, 9]) == ([1, 2, 3, 4, 7, 8, 9], 6)
assert merge_sort([5, 2, 7, 1, 6, 3, 8, 4])[0] == [1, 2, 3, 4, 5, 6, 7, 8]
assert merge_sort([1, 2, 3, 4, 5, 6, 7, 8])[1] == 12
assert merge_sort([8, 7, 6, 5, 4, 3, 2, 1])[1] == 12
assert merge_sort([1, 3, 2, 7, 4, 6, 5, 8])[1] == 17
assert merge_sort([]) == ([], 0)
assert merge_sort([3]) == ([3], 0)
rng = random.Random(1)
for _ in range(200):
sample = [rng.randint(-20, 20) for _ in range(rng.randint(0, 30))]
assert merge_sort(sample)[0] == sorted(sample)
print("ok")퀵 정렬
"""퀵 정렬(로무토 분할): 구간 끝 값을 피벗으로 삼아 작은 값을 왼쪽에 모으고, 피벗을 경계에 놓은 뒤 양쪽을 따로 정렬한다.
시간: 평균 O(n log n), 최악 O(n²) (이미 정렬된 입력·같은 값만 있는 입력 + 맨 끝 피벗)
공간: 재귀 스택 평균 O(log n), 최악 O(n). 제자리 정렬이고 불안정하다.
"""
import random
def partition(values: list, low: int, high: int) -> tuple[int, int]:
"""values[low..high]를 values[high] 기준으로 가르고, 피벗이 선 자리와 비교 횟수를 돌려준다."""
pivot = values[high]
boundary = low
for scan in range(low, high):
if values[scan] < pivot:
values[boundary], values[scan] = values[scan], values[boundary]
boundary += 1
values[boundary], values[high] = values[high], values[boundary]
return boundary, high - low
def quick_sort(values: list, low: int = 0, high: int | None = None) -> int:
"""values를 제자리에서 정렬하고, 피벗과 비교한 횟수를 돌려준다."""
if high is None:
high = len(values) - 1
if low >= high:
return 0
pivot_index, comparisons = partition(values, low, high)
comparisons += quick_sort(values, low, pivot_index - 1)
comparisons += quick_sort(values, pivot_index + 1, high)
return comparisons
def max_depth(values: list, low: int = 0, high: int | None = None, depth: int = 0) -> int:
"""정렬하는 동안 재귀가 내려간 가장 깊은 깊이(맨 위가 0)."""
if high is None:
high = len(values) - 1
if low >= high:
return depth
pivot_index, _ = partition(values, low, high)
return max(
max_depth(values, low, pivot_index - 1, depth + 1),
max_depth(values, pivot_index + 1, high, depth + 1),
)
if __name__ == "__main__":
data = [7, 2, 9, 4, 3, 6]
assert partition(data, 0, len(data) - 1) == (3, 5)
assert data == [2, 4, 3, 6, 9, 7]
data = [5, 3, 8, 1, 9, 2, 7]
quick_sort(data)
assert data == [1, 2, 3, 5, 7, 8, 9]
# 맨 끝 피벗의 최악: 정렬된 입력과 같은 값만 있는 입력
assert quick_sort([1, 2, 3, 4, 5, 6]) == 15
assert max_depth([1, 2, 3, 4, 5, 6]) == 5
assert quick_sort([3, 3, 3, 3, 3]) == 10
assert quick_sort([]) == 0
assert quick_sort([4]) == 0
rng = random.Random(2)
for _ in range(200):
sample = [rng.randint(0, 50) for _ in range(rng.randint(0, 25))]
copy = list(sample)
quick_sort(copy)
assert copy == sorted(sample)
print("ok")힙 정렬
"""힙 정렬: 배열을 최대 힙으로 만든 뒤, 맨 위(최댓값)를 끝으로 보내고 힙을 한 칸 줄이기를 되풀이한다.
시간: 최선·평균·최악 모두 O(n log n) (힙 만들기는 O(n))
공간: O(1) — 배열 안에서 바꾼다. 불안정하다.
"""
import heapq
import random
def sift_down(values: list, root: int, size: int) -> None:
"""values[root]를 더 큰 자식과 바꾸며 내려 values[:size]의 root 아래를 최대 힙으로 맞춘다."""
parent = root
while 2 * parent + 1 < size:
child = 2 * parent + 1
if child + 1 < size and values[child + 1] > values[child]:
child += 1
if values[parent] >= values[child]:
return
values[parent], values[child] = values[child], values[parent]
parent = child
def build_max_heap(values: list) -> None:
"""자식이 있는 마지막 칸부터 거꾸로 올라가며 values 전체를 최대 힙으로 만든다."""
for root in range(len(values) // 2 - 1, -1, -1):
sift_down(values, root, len(values))
def heap_sort(values: list) -> None:
"""values를 제자리에서 오름차순으로 정렬한다."""
build_max_heap(values)
for end in range(len(values) - 1, 0, -1):
values[0], values[end] = values[end], values[0]
sift_down(values, 0, end)
def is_max_heap(values: list) -> bool:
"""모든 부모가 자식보다 크거나 같은지."""
return all(values[(child - 1) // 2] >= values[child] for child in range(1, len(values)))
def top_k(stream: list, k: int) -> list:
"""크기 k인 최소 힙으로 가장 큰 k개를 남겨 오름차순으로 돌려준다. O(n log k)."""
heap: list = []
for value in stream:
if len(heap) < k:
heapq.heappush(heap, value)
elif value > heap[0]:
# 맨 위는 상위 k개에 들기 위한 문턱이다
heapq.heapreplace(heap, value)
return sorted(heap)
if __name__ == "__main__":
assert is_max_heap([9, 7, 8, 3, 6, 5, 2])
assert not is_max_heap([9, 8, 5, 6, 7, 6, 4])
data = [4, 10, 3, 5, 1, 8, 7]
build_max_heap(data)
assert data == [10, 5, 8, 4, 1, 3, 7]
heap_sort(data)
assert data == [1, 3, 4, 5, 7, 8, 10]
assert top_k([5, 12, 3, 9, 15, 7, 11], 3) == [11, 12, 15]
assert top_k([2, 1], 3) == [1, 2]
empty: list = []
heap_sort(empty)
assert empty == []
rng = random.Random(3)
for _ in range(200):
sample = [rng.randint(0, 30) for _ in range(rng.randint(0, 25))]
copy = list(sample)
heap_sort(copy)
assert copy == sorted(sample)
assert top_k(sample, 4) == sorted(sample)[-4:]
print("ok")계수·기수 정렬
"""계수 정렬과 기수 정렬(LSD): 값을 비교하지 않고, 값이나 자리 숫자를 칸 번호로 써서 정렬한다.
시간: 계수 정렬 O(n + k) (k는 값의 범위), 기수 정렬 O(d(n + b)) (d는 자릿수, b는 진법)
공간: 계수 정렬 O(n + k), 기수 정렬 O(n + b). 둘 다 안정 정렬이다.
"""
import random
def counting_sort(values: list[int], max_value: int, key=None) -> list[int]:
"""0~max_value인 키로 안정 정렬한 새 리스트를 돌려준다. key가 없으면 값 자체가 키다."""
key = key or (lambda value: value)
counts = [0] * (max_value + 1)
for value in values:
counts[key(value)] += 1
# 누적합: 키마다 결과에서 시작할 자리
starts = []
total = 0
for count in counts:
starts.append(total)
total += count
output = [0] * len(values)
for value in values:
output[starts[key(value)]] = value
starts[key(value)] += 1
return output
def radix_sort(values: list[int], base: int = 10) -> list[int]:
"""음이 아닌 정수를 낮은 자리부터 자리 숫자로 안정 정렬해 정렬된 새 리스트를 돌려준다."""
result = list(values)
place = 1
largest = max(values, default=0)
while largest // place > 0:
result = counting_sort(result, base - 1, key=lambda value: value // place % base)
place *= base
return result
if __name__ == "__main__":
assert counting_sort([3, 1, 4, 1, 0, 3, 1], 4) == [0, 1, 1, 1, 3, 3, 4]
assert radix_sort([329, 457, 657, 839, 436, 720, 355]) == [329, 355, 436, 457, 657, 720, 839]
# 높은 자리를 먼저 하면 틀린다: 마지막에 쓴 키가 가장 중요한 키가 된다
tens_first = counting_sort([21, 12, 22, 11], 9, key=lambda value: value // 10 % 10)
assert counting_sort(tens_first, 9, key=lambda value: value % 10) == [11, 21, 12, 22]
assert radix_sort([21, 12, 22, 11]) == [11, 12, 21, 22]
# 키가 같으면 들어온 순서를 지킨다(안정)
words = ["b2", "a1", "c2", "d1"]
assert counting_sort(words, 9, key=lambda word: int(word[1])) == ["a1", "d1", "b2", "c2"]
assert counting_sort([], 5) == []
assert radix_sort([]) == []
assert radix_sort([0, 0]) == [0, 0]
rng = random.Random(4)
for _ in range(200):
sample = [rng.randint(0, 9999) for _ in range(rng.randint(0, 30))]
assert radix_sort(sample) == sorted(sample)
assert radix_sort(sample, base=256) == sorted(sample)
small = [value % 50 for value in sample]
assert counting_sort(small, 49) == sorted(small)
print("ok")안정성
"""정렬의 안정성: 키가 같은 값이 정렬 뒤에도 입력 순서를 지키는가.
시간: 예제의 두 정렬은 O(n²), 여러 키 정렬은 키마다 O(n log n) (Python의 안정 정렬 sorted)
공간: O(n) — 입력을 복사해 정렬한다.
"""
import random
def selection_sort_by_key(records: list, key) -> list:
"""멀리 있는 칸과 바꾸는 선택 정렬. 같은 키의 순서가 뒤집힐 수 있다(불안정)."""
result = list(records)
for start in range(len(result) - 1):
min_index = start
for index in range(start + 1, len(result)):
if key(result[index]) < key(result[min_index]):
min_index = index
result[start], result[min_index] = result[min_index], result[start]
return result
def insertion_sort_by_key(records: list, key) -> list:
"""이웃끼리만 옮기는 삽입 정렬. 같은 키를 건너뛰지 않아 안정적이다."""
result = list(records)
for next_index in range(1, len(result)):
item = result[next_index]
hole = next_index
while hole > 0 and key(result[hole - 1]) > key(item):
result[hole] = result[hole - 1]
hole -= 1
result[hole] = item
return result
def is_stable_result(records: list, result: list, key) -> bool:
"""result에서 키가 같은 값끼리 records에서의 순서를 지켰는지."""
position = {id(record): index for index, record in enumerate(records)}
return all(
position[id(left)] < position[id(right)]
for left, right in zip(result, result[1:])
if key(left) == key(right)
)
def number(record: tuple) -> int:
"""레코드의 첫 칸(정렬 키)."""
return record[0]
def sort_by_two_keys(records: list, first_key, second_key) -> list:
"""first_key 순, 같으면 second_key 순. 덜 중요한 키로 먼저, 중요한 키로 나중에 안정 정렬한다."""
return sorted(sorted(records, key=second_key), key=first_key)
if __name__ == "__main__":
cards = [(3, "a"), (1, "b"), (3, "c"), (2, "d"), (1, "e")]
stable = insertion_sort_by_key(cards, number)
unstable = selection_sort_by_key(cards, number)
assert [f"{n}{s}" for n, s in stable] == ["1b", "1e", "2d", "3a", "3c"]
assert [f"{n}{s}" for n, s in unstable] == ["1b", "1e", "2d", "3c", "3a"]
assert is_stable_result(cards, stable, number)
assert not is_stable_result(cards, unstable, number)
trio = [(2, "A"), (2, "B"), (1, "C")]
assert selection_sort_by_key(trio, number) == [(1, "C"), (2, "B"), (2, "A")]
fruits = [(2, "사과"), (1, "귤"), (2, "배"), (1, "감"), (3, "자두"), (2, "포도")]
by_price_then_name = sort_by_two_keys(fruits, number, lambda fruit: fruit[1])
assert by_price_then_name == [(1, "감"), (1, "귤"), (2, "배"), (2, "사과"), (2, "포도"), (3, "자두")]
assert by_price_then_name == sorted(fruits)
assert insertion_sort_by_key([], number) == []
rng = random.Random(5)
for _ in range(200):
sample = [(rng.randint(0, 3), index) for index in range(rng.randint(0, 15))]
result = insertion_sort_by_key(sample, number)
# 안정 정렬이면 Python의 sorted(안정)와 결과가 똑같다
assert result == sorted(sample, key=number)
assert is_stable_result(sample, result, number)
print("ok")이진 탐색과 lower/upper bound
"""이진 탐색과 lower/upper bound: 반열린 구간 [lo, hi)의 가운데 하나만 보고 절반을 버린다.
시간: O(log n), 답을 이진 탐색하는 경우 O(판정 비용 × log(범위))
공간: O(1)
"""
import bisect
import random
def lower_bound(values: list, target) -> int:
"""target 이상인 첫 칸. 없으면 len(values)."""
lo, hi = 0, len(values)
while lo < hi:
mid = lo + (hi - lo) // 2
if values[mid] < target:
lo = mid + 1
else:
hi = mid
return lo
def upper_bound(values: list, target) -> int:
"""target보다 큰 첫 칸. 없으면 len(values)."""
lo, hi = 0, len(values)
while lo < hi:
mid = lo + (hi - lo) // 2
if values[mid] <= target:
lo = mid + 1
else:
hi = mid
return lo
def days_needed(weights: list[int], capacity: int) -> int:
"""짐을 순서대로, 하루 용량 capacity로 나를 때 걸리는 날수."""
days, load = 1, 0
for weight in weights:
if load + weight > capacity:
days += 1
load = 0
load += weight
return days
def min_capacity(weights: list[int], days: int) -> int:
"""days일 안에 나를 수 있는 최소 용량. 판정이 거짓에서 참으로 한 번만 바뀌는 경계를 찾는다."""
lo, hi = max(weights), sum(weights)
while lo < hi:
mid = lo + (hi - lo) // 2
if days_needed(weights, mid) <= days:
hi = mid
else:
lo = mid + 1
return lo
if __name__ == "__main__":
values = [2, 4, 4, 4, 7, 9]
assert [lower_bound(values, 4), upper_bound(values, 4)] == [1, 4]
assert lower_bound(values, 5) == upper_bound(values, 5) == 4
assert upper_bound(values, 10) == 6
assert lower_bound(values, 1) == 0
sorted_values = [1, 3, 4, 7, 7, 9, 12, 15]
assert upper_bound(sorted_values, 7) - lower_bound(sorted_values, 7) == 2
assert lower_bound(sorted_values, 8) == upper_bound(sorted_values, 8) == 5
assert min_capacity([3, 2, 2, 4, 1, 4], 3) == 6
assert min_capacity([5], 1) == 5
assert lower_bound([], 3) == upper_bound([], 3) == 0
rng = random.Random(6)
for _ in range(300):
sample = sorted(rng.randint(0, 20) for _ in range(rng.randint(0, 20)))
target = rng.randint(-1, 21)
assert lower_bound(sample, target) == bisect.bisect_left(sample, target)
assert upper_bound(sample, target) == bisect.bisect_right(sample, target)
print("ok")