def test_initialization(self):
"""Test public methods on an empty union."""
# Create { }.
an_empty_graph = UnionFind()
self.assertEqual(an_empty_graph.n_vertices(), 0)
self.assertEqual(an_empty_graph.connected(0, 0), False)
self.assertEqual(an_empty_graph.connected(0, 1), False)
def test_root(self):
"""Test root()."""
# Create { {0}, {1, 2} }.
a_graph = UnionFind()
a_graph.connect(1, 2)
# Check legal input.
self.assertEqual(a_graph.root(0), 0)
self.assertTrue(a_graph.root(1) in {1, 2})
self.assertTrue(a_graph.root(2) in {1, 2})
# Check illegal input.
with self.assertRaises(AssertionError):
a_graph.root(3)
def test_implicit_adding_by_connecting(self):
"""Test connect(), which implicitly calls add()."""
# Create { }.
a_graph = UnionFind()
# Implicitly add 0, 1, 2 to { }, then connect 1 with 2,
# which makes { {0}, {1, 2} }.
a_graph.connect(1, 2)
self.assertEqual(a_graph.n_vertices(), 3)
# Trivially connected components.
self.assertEqual(a_graph.connected(0, 0), True)
self.assertEqual(a_graph.connected(1, 1), True)
self.assertEqual(a_graph.connected(2, 2), True)
# Connected components created by connect().
self.assertEqual(a_graph.connected(1, 2), True)
self.assertEqual(a_graph.connected(2, 1), True)
# Disconnected components.
self.assertEqual(a_graph.connected(0, 1), False)
self.assertEqual(a_graph.connected(1, 0), False)
self.assertEqual(a_graph.connected(0, 2), False)
self.assertEqual(a_graph.connected(2, 0), False)
def test_non_consecutive_adding(self):
"""Test adding vertices non-consecutively."""
# Create { }.
a_graph = UnionFind()
# Add 1 to { }, which becomes { {0}, {1} }.
# Here 0 is added implicitly.
a_graph.add(1)
self.assertEqual(a_graph.n_vertices(), 2)
self.assertEqual(a_graph.connected(0, 0), True)
self.assertEqual(a_graph.connected(0, 1), False)
self.assertEqual(a_graph.connected(1, 1), True)
def test_consecutive_adding(self):
"""Test adding vertices consecutively."""
# Create { }.
a_graph = UnionFind()
# Add 0 to { }, which becomes { {0} }.
a_graph.add(0)
self.assertEqual(a_graph.n_vertices(), 1)
self.assertEqual(a_graph.connected(0, 0), True)
self.assertEqual(a_graph.connected(0, 1), False)
self.assertEqual(a_graph.connected(1, 1), False)
# Add 0 to { {0} }, nothing changes.
a_graph.add(0)
self.assertEqual(a_graph.n_vertices(), 1)
self.assertEqual(a_graph.connected(0, 0), True)
self.assertEqual(a_graph.connected(0, 1), False)
self.assertEqual(a_graph.connected(1, 1), False)
# Add 1 to { {0} }, which becomes { {0}, {1} }.
a_graph.add(1)
self.assertEqual(a_graph.n_vertices(), 2)
self.assertEqual(a_graph.connected(0, 0), True)
self.assertEqual(a_graph.connected(0, 1), False)
self.assertEqual(a_graph.connected(1, 1), True)
class Scheduler:
"""A scheduler supporting O(1) adding and O(N) scheduling."""
def __init__(self):
self._task_to_id = dict()
self._id_to_task = list()
self._graph = DirectedGraph()
self._union = UnionFind()
def add_a_task(self, task):
"""Add a new task.
Do nothing, if the task has already been added.
"""
if task not in self._task_to_id:
i_task = self.n_tasks()
self._task_to_id[task] = i_task
self._id_to_task.append(task)
assert len(self._id_to_task) == len(self._task_to_id)
assert task == self._id_to_task[self._task_to_id[task]]
def add_tasks(self, tasks):
"""Add multiple tasks."""
for task in tasks:
self.add_a_task(task)
def n_tasks(self):
"""Return the number of tasks being added."""
return len(self._task_to_id)
def add_a_prerequisite(self, task, prerequisite):
"""Add a prerequisite for a task.
Automatically add a new task, if any of the two is new.
Do nothing, if the prerequisite has already been added.
"""
self.add_a_task(task)
self.add_a_task(prerequisite)
i_task = self._task_to_id[task]
i_prerequisite = self._task_to_id[prerequisite]
self._graph.connect(i_task, i_prerequisite)
self._union.connect(i_task, i_prerequisite)
def add_prerequisites(self, task, prerequisites):
"""Add multiple prerequisites for a task."""
for prerequisite in prerequisites:
self.add_a_prerequisite(task, prerequisite)
def schedule(self):
"""Return the tasks in topologically sorted order."""
sorted_tasks = TopologicalSort(self._graph).sort()
# Make immutable copies.
scheduled_tasks = set()
for a_component in self._to_components(sorted_tasks):
scheduled_tasks.add(tuple(a_component))
return scheduled_tasks
def _to_components(self, sorted_tasks):
root_to_component = dict()
for i_task in sorted_tasks:
i_root = self._union.root(i_task)
if i_root not in root_to_component:
root_to_component[i_root] = list()
task = self._id_to_task[i_task]
root_to_component[i_root].append(task)
return root_to_component.values()
def __init__(self):
self._task_to_id = dict()
self._id_to_task = list()
self._graph = DirectedGraph()
self._union = UnionFind()