Templates

1# this template is typically used to answer question such as
2# find **all** the combination of ...
3# generate **all** valid...
4ans = []
5def dfs(start_index, path, [...additional states]):
6    if is_leaf(start_index):
7        ans.append(path[:]) # add a copy of the path to the result
8        return
9    for edge in get_edges(start_index, [...additional states]):
10        # prune if needed
11        if not is_valid(edge):
12            continue
13        path.add(edge)
14        if additional states:
15            update(...additional states)
16        dfs(start_index + len(edge), path, [...additional states])
17        # revert(...additional states) if necessary e.g. permutations
18        path.pop()
19

1def binary_search(arr: List[int], target: int) -> int:
2    left, right = 0, len(arr) - 1
3    first_true_index = -1
4    while left <= right:
5        mid = (left + right) // 2
6        if feasible(mid):
7            first_true_index = mid
8            right = mid - 1
9        else:
10            left = mid + 1
11
12    return first_true_index
13

1def bfs(root):
2    queue = deque([root])
3    while len(queue) > 0:
4        node = queue.popleft()
5        for child in node.children:
6            if is_goal(child):
7                return FOUND(child)
8            queue.append(child)
9    return NOT_FOUND
10

1def dfs(root, target):
2    if root is None:
3        return None
4    if root.val == target:
5        return root
6    left = dfs(root.left, target)
7    if left is not None:
8        return left
9    return dfs(root.right, target)
10

1def bfs(root):
2    queue = deque([root])
3    visited = set([root])
4    while len(queue) > 0:
5        node = queue.popleft()
6        for neighbor in get_neighbors(node):
7            if neighbor in visited:
8                continue
9            queue.append(neighbor)
10            visited.add(neighbor)
11

1def dfs(root, visited):
2    for neighbor in get_neighbors(root):
3        if neighbor in visited:
4            continue
5        visited.add(neighbor)
6        dfs(neighbor, visited)
7

1num_rows, num_cols = len(grid), len(grid[0])
2def get_neighbors(coord):
3    row, col = coord
4    delta_row = [-1, 0, 1, 0]
5    delta_col = [0, 1, 0, -1]
6    res = []
7    for i in range(len(delta_row)):
8        neighbor_row = row + delta_row[i]
9        neighbor_col = col + delta_col[i]
10        if 0 <= neighbor_row < num_rows and 0 <= neighbor_col < num_cols:
11            res.append((neighbor_row, neighbor_col))
12    return res
13
14from collections import deque
15
16def bfs(starting_node):
17    queue = deque([starting_node])
18    visited = set([starting_node])
19    while len(queue) > 0:
20        node = queue.popleft()
21        for neighbor in get_neighbors(node):
22            if neighbor in visited:
23                continue
24            # Do stuff with the node if required
25            # ...
26            queue.append(neighbor)
27            visited.add(neighbor)
28

1def mono_stack(insert_entries):
2    stack = []
3    for entry in insert_entries:
4        while stack and stack[-1] <= entry:
5            stack.pop()
6            # Do something with the popped item here
7        stack.append(entry)
8

1def count_parents(graph):
2    counts = { node: 0 for node in graph }
3    for parent in graph:
4        for node in graph[parent]:
5            counts[node] += 1
6    return counts
7
8def topo_sort(graph):
9    res = []
10    q = deque()
11    counts = count_parents(graph)
12    for node in counts:
13        if counts[node] == 0:
14            q.append(node)
15    while len(q) > 0:
16        node = q.popleft()
17        res.append(node)
18        for child in graph[node]:
19            counts[child] -= 1
20            if counts[child] == 0:
21                q.append(child)
22    return res
23

1class Node:
2    def __init__(self, value):
3        self.value = value
4        self.children = {}
5
6    def insert(self, s, idx):
7        # idx: index of the current character in s
8        if idx != len(s):
9            self.children.setdefault(s[idx], Node(s[idx]))
10            self.children.get(s[idx]).insert(s, idx + 1)
11

1class UnionFind:
2    def __init__(self):
3        self.id = {}
4
5    def find(self, x):
6        y = self.id.get(x, x)
7        if y != x:
8            self.id[x] = y = self.find(y)
9        return y
10
11    def union(self, x, y):
12        self.id[self.find(x)] = self.find(y)
13