Graph Algorithms - CSES Problem Set Solutions

I solved all problems on this page without any hints/spoilers (not even reading the CSES recommended book).

Visit https://cses.fi/problemset/ for the full problem set.

I don't provide the full problem specifications on this page due to possible copyright issues. Links to the original problem specs are provided below along with the date accessed, which should allow you to use Internet Archive if the original URL hosting a problem specification ever meaningfully changes.

Counting Rooms [Spec] (2024-03-13)

My solution involves iterating through every cell of the graph, and if a room is discovered, we flood fill the room with '#' characters and increment a counter by 1.

My first attempt involved a recursive flood fill, which caused segfaults due to stack overflow:

Switching the flood fill implementation to an iterative one passed all test cases when using PyPy3 (but TLE’s when using CPython3):

#!/usr/bin/env python3

from sys import stdin, setrecursionlimit

setrecursionlimit(10000000)

stdin.readline()
graph = [[*s.strip(), '#'] for s in stdin.readlines()]

graph.append(['#']*len(graph[0]))
(len1, len2) = (len(graph), len(graph[0]))

def floodfill(i, j):
    stack = [(i, j)]
    while len(stack):
        (i, j) = stack.pop()
        if graph[i][j] == '#':
            continue
        graph[i][j] = '#'
        stack.extend(((i + 1, j), (i - 1, j), (i, j + 1), (i, j - 1)))

solution = 0
for i in range(len1):
    for j in range(len2):
        if graph[i][j] == '.':
            solution += 1
            floodfill(i, j)
print(solution)

Labyrinth [Spec] (2024-03-16)

This took a very unnecessarily long path towards reaching a successful Python solution that started with a naive BFS solution. That TLE’d, so I tried multiple A* variants which also TLE’d. When I circled back around to trying BFS with a simple optimization that reduced overhead, which passed.

TLE’d solutions in the order I tried them:

When I applied the ‘seen set’ optimization to my original BFS solution, it passed all tests in PyPy3 (but TLE’s in CPython3):

#!/usr/bin/env python3

from sys import stdin
from collections import deque

stdin.readline()
graph = [[*s.strip(), '#'] for s in stdin.readlines()]

graph.append(['#']*len(graph[0]))
(len1, len2) = (len(graph), len(graph[0]))

a_loc = None
b_loc = None
for i in range(len1):
    for j in range(len2):
        if graph[i][j] == 'A':
            a_loc = (i, j)
            graph[i][j] = '#'
        elif graph[i][j] == 'B':
            b_loc = (i, j)
            graph[i][j] = '.'

dq = deque([(*a_loc, None)])
solution = None
while (solution is None) and len(dq):
    tup = dq.popleft()
    (i, j, _) = tup
    for ii, jj in ((i + 1, j), (i - 1, j), (i, j + 1), (i, j - 1)):
        if graph[ii][jj] == '#':
            continue
        elif (ii, jj) == b_loc:
            solution = (ii, jj, tup)
            break
        graph[ii][jj] = '#'
        dq.append((ii, jj, tup))

if solution is None:
    print("NO")
else:
    path = []
    while solution[2] is not None:
        prev = solution[2]
        (difi, difj) = (solution[0] - prev[0], solution[1] - prev[1])
        path.append(
            'D' if (difi == 1) else
            'U' if (difi == -1) else
            'R' if (difj == 1) else
            'L'
        )
        solution = prev
    print("YES")
    print(len(path))
    print("".join(reversed(path)))

Building Roads [Spec] (2024-03-18)

I used a DFS approach, first traversing every node reachable from node 1 and marking them all as seen. After the DFS operation concludes and there are still unseen nodes, we build a road from node 1 to any unseen node, then we DFS from that unseen node. We repeat until all nodes have been seen.

I initially implemented DFS using recursion, but ran into stack overflow errors:

Reimplementing DFS using a stack data structure passed all tests:

#!/usr/bin/env python3

from sys import stdin
from collections import defaultdict

(n, _) = [int(s) for s in stdin.readline().split()]

graph = defaultdict(list)
for s in stdin.readlines():
    (a, b) = tuple(int(x) for x in s.split())
    graph[a].append(b)
    graph[b].append(a)

unseen = set(range(2, n + 1))

def dfs(i):
    stack = [i]
    while len(stack):
        i = stack[-1]
        if len(graph[i]):
            j = graph[i].pop()
            if j in unseen:
                unseen.remove(j)
                stack.append(j)
        else:
            stack.pop()

solution = []

dfs(1)
while len(unseen):
    i = unseen.pop()
    solution.append((1, i))
    dfs(i)

print(len(solution))
print("\n".join(f"{a} {b}" for a, b in solution))

I also implemented a simple graph representation optimization, which did improve the performance slightly:

I also tried union find to replace the unseen set, but it ran slightly slower:

Message Route [Spec] (2024-03-18)

A very simple BFS solution:

#!/usr/bin/env python3

from sys import stdin
from collections import deque

(n, _) = [int(s) for s in stdin.readline().split()]

graph = [[] for _ in range(n + 1)]  # 0th index isn't used
for s in stdin.readlines():
    (a, b) = tuple(int(x) for x in s.split())
    graph[a].append(b)
    graph[b].append(a)

seen = {}  # {node: previous node before it}

dq = deque([1])
solution = None
while (not solution) and len(dq):
    i = dq.popleft()
    for j in graph[i]:
        if j in seen:
            continue
        elif j == n:
            solution = i
            break
        seen[j] = i
        dq.append(j)

if solution is None:
    print("IMPOSSIBLE")
else:
    path = [n]
    while solution != 1:
        path.append(solution)
        solution = seen[solution]
    print(len(path) + 1)
    print("1 " + " ".join(str(x) for x in reversed(path)))

Building Teams [Spec] (2024-03-19)

The idea is to BFS, and assign one layer of nodes at a time. When we check through a layer’s connections, if we see that there is a connection within the layer, we terminate the algorithm early and print “IMPOSSIBLE”. Full solution:

#!/usr/bin/env python3

from sys import stdin

(n, _) = [int(s) for s in stdin.readline().split()]

graph = [[] for _ in range(n + 1)]  # 0th index isn't used
for s in stdin.readlines():
    (a, b) = tuple(int(x) for x in s.split())
    graph[a].append(b)
    graph[b].append(a)

teams = [0 for _ in range(n + 1)]  # 0th index isn't used

def bfs(start):
    teams[start] = 1
    (cur, cur_team) = ([start], 1)
    while len(cur):
        nxt = []
        nxt_team = 2 if (cur_team == 1) else 1
        for i in cur:
            for j in graph[i]:
                if teams[j] == cur_team:
                    return False
                elif not teams[j]:
                    teams[j] = nxt_team
                    nxt.append(j)
        (cur, cur_team) = (nxt, nxt_team)
    return True

impossible = False
for i in range(1, n + 1):
    if (not teams[i]) and (not bfs(i)):
        impossible = True
        break
print("IMPOSSIBLE" if impossible else " ".join(str(x) for x in teams[1:]))

Round Trip [Spec] (2024-03-20)

I went for a DFS/backtracking approach. We basically attempt every possible path, going as deep as possible and if we find we can loop back to another node along the current path, we terminate.

Full implementation:

#!/usr/bin/env python3

from sys import stdin, setrecursionlimit
setrecursionlimit(1000000)

(n, _) = [int(s) for s in stdin.readline().split()]

graph = [[] for _ in range(n + 1)]  # 0th index isn't used
for s in stdin.readlines():
    (a, b) = tuple(int(x) for x in s.split())
    graph[a].append(b)
    graph[b].append(a)

level = [0 for _ in range(n + 1)]     # 0th index isn't used
seen = [False for _ in range(n + 1)]  # 0th index isn't used

def backtrack(cur, cur_level):
    seen[cur] = True
    level[cur] = cur_level
    for nxt in graph[cur]:
        nxt_level = level[nxt]
        if (nxt_level > 0) and (cur_level > nxt_level + 1):
            return [nxt, cur]
        elif level[nxt]:
            continue
        v = backtrack(nxt, cur_level + 1)
        if v:
            if v[0] != v[-1]:
                v.append(cur)
            return v
    level[cur] = 0
    return None

sol = None
for i in range(1, n + 1):
    if not seen[i]:
        sol = backtrack(i, 1)
        if sol:
            break

print((f"{len(sol)}\n" + " ".join(str(x) for x in sol)) if sol else "IMPOSSIBLE")

In theory, if the entire graph is fully connected (i.e. there is always a path between any two nodes in the entire graph), we can always start our algorithm at any arbitrary node, and that would be enough to solve the problem. However, we find that there are test cases involving multiple components (i.e. you can find two nodes in the graph such that there is no path between them). To test this, I added a single break (the third-last line), which failed some tests (as predicted if the graph has multiple components):

Multiple components is the reason why the seen set (implemented as a list) and the for i in range(1, n + 1): loop are included in my solution. The seen set and that loop lets my code scan for components that haven’t been dealt with yet before terminating.

Interestingly, I switch the line elif level[nxt]: to elif seen[nxt]: and it passes all the tests, but I’m still not entirely sure why.

TODO: Find out why!

Monsters [Spec] (2024-03-20)

I used a BFS solution that deals with one layer at a time.

The main loop starts by propagating monster positions by one layer, blocking off squares by writing '#' characters to the graph. Where the monsters are, they may as well be walls. The second part within the main loop propagates player positions by one layer.

For convenience, I add a new column to the right-most and a new row at the bottom-most, all filled with 'X'. This character represents the boundary of the graph. This works not just when our search algorithm attempts to move past the right edge or the bottom edge, but also when the algo attempts to move past the left edge or the top edge. This is because of Python’s negative indexing, meaning an index of -1 will reference the end of the list. With the boundary indicated by 'X' characters, we no longer need to explicitly check if an index exists within the graph, or if the player has reached the boundary.

#!/usr/bin/env python3

from sys import stdin
from collections import deque

# I'm adding character 'X' to indicate the outside of the graph.

stdin.readline()
graph = [[*s.strip(), 'X'] for s in stdin.readlines()]

graph.append(['X']*len(graph[0]))
(len1, len2) = (len(graph), len(graph[0]))

player_dq = deque()
monsters_dq = deque()

for i in range(len1):
    for j in range(len2):
        if graph[i][j] == 'M':
            monsters_dq.append((i, j))
            graph[i][j] = '#'
        elif graph[i][j] == 'A':
            player_dq.append((i, j, None))
            graph[i][j] = '#'

directions = ((1, 0), (-1, 0), (0, 1), (0, -1))

solution = None
while (solution is None) and len(player_dq):
    # propagate monster squares
    for _ in range(len(monsters_dq)):
        (i, j) = monsters_dq.popleft()
        for di, dj in directions:
            (ii, jj) = (i + di, j + dj)
            if graph[ii][jj] == '.':
                graph[ii][jj] = '#'
                monsters_dq.append((ii, jj))
    # propagate player positions
    for _ in range(len(player_dq)):
        tup = player_dq.popleft()
        (i, j, prev) = tup
        for di, dj in directions:
            (ii, jj) = (i + di, j + dj)
            if graph[ii][jj] == 'X':
                solution = tup
                break
            elif graph[ii][jj] == '.':
                graph[ii][jj] = '#'
                player_dq.append((ii, jj, tup))
        if solution:
            break

if solution is None:
    print("NO")
else:
    path = []
    while solution[2] is not None:
        prev = solution[2]
        (difi, difj) = (solution[0] - prev[0], solution[1] - prev[1])
        path.append(
            'D' if (difi == 1) else
            'U' if (difi == -1) else
            'R' if (difj == 1) else
            'L'
        )
        solution = prev
    print("YES")
    print(len(path))
    print("".join(reversed(path)))

Shortest Routes I [Spec] (2024-03-20)

I used a simple modified Dijkstra’s algorithm that passes through all nodes. Note that every time we pop from the priority queue, it is guaranteed that the distance to that node is the shortest distance to that node, so for convenience, we can prune the search by ignoring anything where we already know the shortest path to.

#!/usr/bin/env python3

from sys import stdin
from heapq import heappush, heappop

(n, _) = [int(s) for s in stdin.readline().split()]

graph = [[] for _ in range(n + 1)]  # 0th index isn't used
for s in stdin.readlines():
    (a, b, w) = tuple(int(x) for x in s.split())
    graph[a].append((b, w))

dists = [-1 for _ in range(n + 1)]  # 0th index isn't used

pq = [(0, 1)]
while len(pq):
    (dist, i) = heappop(pq)
    if dists[i] == -1:
        dists[i] = dist
        for j, w in graph[i]:
            heappush(pq, (dist + w, j))

print(" ".join(str(x) for x in dists[1:]))