코딩 테스트 유형별 소스 총 정리

BFS

def bfs(graph, visited, start):
    q = deque([start])
    visited[start] = True

    while q:
        cur = q.popleft()

        for i in range(n):
            if visited[cur]:
                continue

            visit[i] = True
            q.append(i)

DFS

def dfs(graph, visited, start):
    visited[start] = True

    for nxt in graph[start]:
        if not visited[nxt]:
            dfs(nxt)

백트래킹, 정규표현식

프로그래머스: 2019 카카오 개발자 겨울 인턴십 - 불량 사용자

import re
from copy import deepcopy

answer = []
banned_id_matches = []


def solution(user_ids, banned_ids):
    global answer, banned_id_matches

    set_banned_id_matches(user_ids, banned_ids)
    search(set(), 0)

    return len(answer)


def set_banned_id_matches(user_ids, banned_ids):
    for i in range(len(banned_ids)):
        p = re.compile("^" + banned_ids[i].replace("*", ".") + "$")
        banned_id_matches.append({user for user in user_ids if p.match(user)})


def search(visited, idx):
    if idx == len(banned_id_matches):
        if visited not in answer:
            answer.append(deepcopy(visited))
        return

    for user in banned_id_matches[idx]:
        if user in visited:
            continue

        visited.add(user)
        search(visited, idx + 1)
        visited.remove(user)

정규표현식

프로그래머스: 2019 KAKAO BLIND RECRUITMENT - 매칭 점수

import re
from collections import defaultdict

link_relations = defaultdict(list)
link_dict = {}


def solution(word, pages):
    result = []

    my_link_re = re.compile('<meta(.+?)/>', re.MULTILINE)
    links_re = re.compile('<a(.+?)>', re.MULTILINE)
    # words_re = re.compile('[^a-zA-Z]'+word+'[^a-zA-Z]', re.I)

    for idx, page in enumerate(pages):
        my_link = getUrl(my_link_re.search(page).group())
        links = [getUrl(link) for link in links_re.findall(page)]
        score = re.sub('[^a-z]+', '.', page.lower()).split('.').count(word.lower())

        link_dict[my_link] = { 'index': idx, 'score': score, 'link_count': len(links) }

        for link in links:
            link_relations[link].append(my_link)

    for my_link in link_dict.keys():
        idx = link_dict[my_link]['index']
        score = link_dict[my_link]['score']
        matching_score = score + getLinkScore(my_link)

        result.append((matching_score, idx))

    return sorted(result, key = lambda x: (-x[0], x[1]))[0][1]


def getUrl(str):
    return str.split('"')[-2]


def getLinkScore(my_link):
    link_score = 0

    for linked_link in link_relations[my_link]:
        if link_dict[linked_link]['link_count'] == 0:
            continue
        link_score += link_dict[linked_link]['score'] / link_dict[linked_link]['link_count']

    return link_score

이진탐색

def binary_search(arr, target, start, end):
    while start <= end:
        mid = (start + end) // 2

        if arr[mid] == target:
            return mid

        if arr[mid] < target:
            start = mid + 1
        else:
            end = mid - 1

    return None

파라메트릭 서치

프로그래머스: 입국심사

def solution(n, times):
    answer = 0
    start, end, mid = 1, times[-1] * n, 0

    while start < end:
        mid = (start + end) // 2

        total = 0
        for time in times:
            total += mid // time

        if total < n:
            start = mid + 1
        else:
            end = mid

    answer = start

    return answer

프로그래머스: 2019 카카오 개발자 겨울 인턴십 - 징검다리 건너기

def solution(stones, k):
    left, right = 1, 200000001
    # 마지막 사람이 건너기 위해 그 전에 최대로 건널 수 있는 수
    result = 0

    while left <= right:
        mid = (left + right) // 2

        if isNextPeoplePossibleCross(mid, stones, k):
            left = mid + 1
            result = mid
        else:
            right = mid - 1

    # 마지막 사람까지 건너면 1을 더해줘야 함
    return result + 1


def isNextPeoplePossibleCross(num, stones, k):
    count = 0
    for i in range(len(stones)):
        if stones[i] - num <= 0:
            count += 1
        else:
            count = 0

        if count >= k:
            return False

    return True

다익스트라 (a에서 b까지 최단경로)

def solution(n, costs):
    graph = [[] for _ in range(n)]

    for a, b, cost in range(costs):
        graph[a][b] = cost

    dists = [math.inf] * n

    dijkstra(graph, dists, 0)


def dijkstra(graph, dists, start):
    pq = []
    heappush(pq, (0, start))
    dists[start] = 0

    while pq:
        nowDist, now = heappop(pq)

        if dists[now] < nowDist:
            continue

        for nxt, nxtCost in graph[now]:
            dist = nowDist + nxtCost

            if dist < dists[nxt]:
                dists[nxt] = dist
                heappush(pq, (dist, nxt))

플로이드 워셜 (노드들에서 노드들까지 모든 최단경로)

import math

# n은 노드 수
def solution(n, costs):
    graph = [[math.inf] * n for _ in range(n)]

    for a in range(n):
        for b in range(n):
            if a == b:
                graph[a][b] = 0

    for a, b, cost in range(costs):
        graph[a][b] = cost

    for k in range(n):
        for a in range(n):
            for b in range(n):
                if a == k or b == k or a == b:
                    continue
                graph[a][b] = min(graph[a][b], graph[a][k] + graph[k][b])

MST - 프림

from heapq import heappush, heappop

def solution(n, costs):
    graph = [[] for _ in range(101)]

    for a, b, cost in costs:
        graph[a].append((b, cost))
        graph[b].append((a, cost))

    return getMinCost(graph)


def getMinCost(graph):
    visited = set([0])
    pq = []
    result = 0

    for nxtNode, nxtCost in graph[0]:
        heappush(pq, (nxtCost, nxtNode))

    while pq:
        curCost, curNode = heappop(pq)

        if curNode in visited:
            continue

        visited.add(curNode)
        result += curCost

        for nxtNode, nxtCost in graph[curNode]:
            heappush(pq, (nxtCost, nxtNode))

    return result

서로소 집합, Union Find

Find

# 자식 노드들이 모두 최고 조상 부모를 가리킬 경우 (경로 압축)
def find(node_num):
    if parent[node_num] != node_num:
        parent[node_num] = find(parent[node_num])

    return parent[node_num]

# 자식 노드들이 기존 부모를 기억해야되는 경우 (탐색 비효율)
def find(node_num):
    if parent[node_num] != node_num:
        return find(parent[node_num])

    return node_num

Union

# 단순히 부모 노드 숫자가 작은 쪽으로 합치기
def union(node1_num, node2_num):
    node1_root = find(node1_num)
    node2_root = find(node2_num)

    if node1_root == node2_root:
        return

    if rank[node1_root] < rank[node2_root]:
        parent[node2_root] = node1_root
    else:
        parent[node1_root] = node2_root

# 랭크값이 높은 부모 쪽으로 합치기
def union(node1_num, node2_num):
    node1_root = find(node1_num)
    node2_root = find(node2_num)

    if node1_root != node2_root:
        if rank[node1_root] > rank[node2_root]:
            parent[node2_root] = node1_root
        else:
            parent[node1_root] = node2_root

            if rank[node1_root] == rank[node2_root]:
                rank[node2_root] += 1

서로소 집합 사이클 판별

cycle = False

for i in range(n):
    a, b = map(int, input().split())

    if find(parents, a) == find(parents, b):
        cycle = True
        break

    union(parents, a, b)

MST - 크루스칼

parent = {}
rank = {}


def solution(n, costs):
    return kruskal(n, costs)


def kruskal(n, costs):
    for v in range(n):
        make_set(v)

    totalCost = 0
    costs.sort(key=lambda x: x[2])

    for node1_num, node2_num, cost in costs:
        if find(node1_num) != find(node2_num):
            union(node1_num, node2_num)
            totalCost += cost

    return totalCost


def make_set(node_num):
    parent[node_num] = node_num
    # 랭크 필요할 때만
    rank[node_num] = 0

위상정렬

def findOrder(self, numCourses: int, prerequisites: List[List[int]]) -> List[int]:
    g = defaultdict(list)
    indegrees = [0 for _ in range(numCourses)]

    for a, b in prerequisites:
        g[a].append(b)
        indegrees[b] += 1

    q = deque()
    for node in range(numCourses):
        if indegrees[node] == 0:
            q.append(node)

    ret = []
    while q:
        cur = q.popleft()
        ret.append(cur)
        for nxt_node in g[cur]:
            indegrees[nxt_node] -= 1
            if indegrees[nxt_node] == 0:
                q.append(nxt_node)

    return ret[::-1] if len(ret) == numCourses else []

트라이 자료구조

프로그래머스: 2020 KAKAO BLIND RECRUITMENT - 가사 검색

class Node:
    def __init__(self, data):
        self.data = data
        self.count = 0
        self.child = {}


class Trie:
    def __init__(self):
        self.head = Node(None)

    def insert(self, string):
        cur = self.head
        cur.count += 1

        for c in string:
            if c not in cur.child:
                cur.child[c] = Node(c)
            cur = cur.child[c]
            cur.count += 1

    def count(self, prefix):
        cur = self.head

        for c in prefix:
            if c not in cur.child:
                return 0
            cur = cur.child[c]

        return cur.count


def solution(words, queries):
    answer = []

    tries = create_trie(words)
    reversed_tries = create_trie(words, True)

    for query in queries:
        answer.append(count_matched_word(tries, reversed_tries, query))

    return answer


def create_trie(words, is_reversed=False):
    trie_dic = {i: Trie() for i in range(1, 10001)}

    for word in words:
        if is_reversed:
            word = word[::-1]
        trie_dic[len(word)].insert(word)

    return trie_dic


def count_matched_word(tries, reversed_tries, query):
    no_mark_query = query.replace('?', '')

    if query[0] == '?':
        return reversed_tries[len(query)].count(no_mark_query[::-1])
    else:
        return tries[len(query)].count(no_mark_query)

KMP

def kmp(s, p):
    res = []
    pi = get_pi(p)
    n, m, j = len(s), len(p), 0

    for i in range(n):
        while j > 0 and s[i] != p[j]:
            j = pi[j-1]

        if s[i] == p[j]:
            if j == m-1:
                res.append(i-m+1)
                j = pi[j]
            else:
                j += 1

    return res


def get_pi(p):
    m, j = len(p), 0
    pi = [0 for _ in range(m)]

    for i in range(1, m):
        while j > 0 and p[i] != p[j]:
            j = pi[j-1]

        if p[i] == p[j]:
            j += 1
            pi[i] = j

    return pi

LIS

LIS 길이만 구하기 (N^2)

N = 6
arr = [2, 3, 1, 5, 4, 6]
dp = [1] * 6

for i in range(1, N):
    for j in range(i):
        if arr[i] > arr[j]:
            dp[i] = max(dp[i], dp[j] + 1)

LIS 길이만 구하기 (N*logN)

from bisect import bisect_left

arr = [1, 3, 4, 2, 1]
table = []
for num in arr:
    if table[-1] < num:
        table.append(num)
    else:
        idx = bisect_left(table, num)
        table[idx] = num

LIS 역추적 (n^2)

N = 6
arr = [2, 3, 1, 5, 4, 6]
dp = [1] * 6

for i in range(1, N):
    for j in range(i):
        if arr[i] > arr[j]:
            dp[i] = max(dp[i], dp[j] + 1)

order = max(dp)
path = []

for i in range(N - 1, -1, -1):
    if dp[i] == order:
        path.append(arr[i])
        order -= 1

path.reverse()

LCS

def lcs(a, b):
    m = len(a)
    n = len(b) 
    dp = [[None] * (n + 1) for _ in range(m + 1)]

    for i in range(m + 1):
        for j in range(n + 1):
            if i == 0 or j == 0:
                dp[i][j] = 0

            elif X[i - 1] == Y[j - 1]:
                dp[i][j] = dp[i - 1][j - 1] + 1

            else:
                dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])

펜윅 트리 (부분합)

class FenwickTree(object):
    def __init__(self, nums):
        n = len(nums)
        self.nums = [0 for _ in range(n+1)]
        self.N = [0 for _ in range(n+1)]
        for i, v in enumerate(nums):
            self.set(i+1, v)

    def _lowbit(self, a):
        return a & -a

    def set(self, i, val):
        diff = val - self.nums[i]
        self.nums[i] = val
        while i < len(self.N):
            self.N[i] += diff
            i += self._lowbit(i)

    def get(self, i):
        ret = 0
        while i > 0:
            ret += self.N[i]
            i -= self._lowbit(i)

        return ret


class Fenwick(object):
    def __init__(self, nums):
        self.bit = FenwickTree(nums)

    def update(self, i, val):
        self.bit.set(i+1, val)

    def sumRange(self, i, j):
        return self.bit.get(j+1)-self.bit.get(i)

LRU 캐시 (O(1))

from collections import OrderedDict

class LRUCache:
    def __init__(self, Capacity):
        self.size = Capacity
        self.cache = OrderedDict()

    def get(self, key):
        if key not in self.cache: return -1
        val = self.cache[key]
        self.cache.move_to_end(key)
        return val

    def put(self, key, val):
        if key in self.cache: del self.cache[key]
        self.cache[key] = val
        if len(self.cache) > self.size:
            self.cache.popitem(last=False)