diff --git a/miplearn/problems/tsp.py b/miplearn/problems/tsp.py
index 47f304b..7ae58e1 100644
--- a/miplearn/problems/tsp.py
+++ b/miplearn/problems/tsp.py
@@ -1,6 +1,7 @@
 #  MIPLearn: Extensible Framework for Learning-Enhanced Mixed-Integer Optimization
 #  Copyright (C) 2020-2021, UChicago Argonne, LLC. All rights reserved.
 #  Released under the modified BSD license. See COPYING.md for more details.
+from typing import List, Tuple, FrozenSet, Any, Optional
 
 import networkx as nx
 import numpy as np
@@ -11,16 +12,16 @@ from scipy.stats import uniform, randint
 from scipy.stats.distributions import rv_frozen
 
 from miplearn.instance.base import Instance
+from miplearn.types import VariableName, Category
 
 
 class ChallengeA:
     def __init__(
         self,
-        seed=42,
-        n_training_instances=500,
-        n_test_instances=50,
-    ):
-
+        seed: int = 42,
+        n_training_instances: int = 500,
+        n_test_instances: int = 50,
+    ) -> None:
         np.random.seed(seed)
         self.generator = TravelingSalesmanGenerator(
             x=uniform(loc=0.0, scale=1000.0),
@@ -38,33 +39,103 @@ class ChallengeA:
         self.test_instances = self.generator.generate(n_test_instances)
 
 
+class TravelingSalesmanInstance(Instance):
+    """An instance ot the Traveling Salesman Problem.
+
+    Given a list of cities and the distance between each pair of cities, the problem
+    asks for the shortest route starting at the first city, visiting each other city
+    exactly once, then returning to the first city. This problem is a generalization
+    of the Hamiltonian path problem, one of Karp's 21 NP-complete problems.
+    """
+
+    def __init__(self, n_cities: int, distances: np.ndarray) -> None:
+        super().__init__()
+        assert isinstance(distances, np.ndarray)
+        assert distances.shape == (n_cities, n_cities)
+        self.n_cities = n_cities
+        self.distances = distances
+        self.edges = [
+            (i, j) for i in range(self.n_cities) for j in range(i + 1, self.n_cities)
+        ]
+        self.varname_to_index = {f"x[{e}]": e for e in self.edges}
+
+    @overrides
+    def to_model(self) -> pe.ConcreteModel:
+        model = pe.ConcreteModel()
+        model.x = pe.Var(self.edges, domain=pe.Binary)
+        model.obj = pe.Objective(
+            expr=sum(model.x[i, j] * self.distances[i, j] for (i, j) in self.edges),
+            sense=pe.minimize,
+        )
+        model.eq_degree = pe.ConstraintList()
+        model.eq_subtour = pe.ConstraintList()
+        for i in range(self.n_cities):
+            model.eq_degree.add(
+                sum(
+                    model.x[min(i, j), max(i, j)]
+                    for j in range(self.n_cities)
+                    if i != j
+                )
+                == 2
+            )
+        return model
+
+    @overrides
+    def get_variable_category(self, var_name: VariableName) -> Category:
+        return self.varname_to_index[var_name]
+
+    @overrides
+    def find_violated_lazy_constraints(self, model: Any) -> List[FrozenSet]:
+        selected_edges = [e for e in self.edges if model.x[e].value > 0.5]
+        graph = nx.Graph()
+        graph.add_edges_from(selected_edges)
+        components = [frozenset(c) for c in list(nx.connected_components(graph))]
+        violations = []
+        for c in components:
+            if len(c) < self.n_cities:
+                violations += [c]
+        return violations
+
+    @overrides
+    def build_lazy_constraint(self, model: Any, component: FrozenSet) -> Any:
+        cut_edges = [
+            e
+            for e in self.edges
+            if (e[0] in component and e[1] not in component)
+            or (e[0] not in component and e[1] in component)
+        ]
+        return model.eq_subtour.add(sum(model.x[e] for e in cut_edges) >= 2)
+
+
 class TravelingSalesmanGenerator:
     """Random generator for the Traveling Salesman Problem."""
 
     def __init__(
         self,
-        x=uniform(loc=0.0, scale=1000.0),
-        y=uniform(loc=0.0, scale=1000.0),
-        n=randint(low=100, high=101),
-        gamma=uniform(loc=1.0, scale=0.0),
-        fix_cities=True,
-        round=True,
-    ):
+        x: rv_frozen = uniform(loc=0.0, scale=1000.0),
+        y: rv_frozen = uniform(loc=0.0, scale=1000.0),
+        n: rv_frozen = randint(low=100, high=101),
+        gamma: rv_frozen = uniform(loc=1.0, scale=0.0),
+        fix_cities: bool = True,
+        round: bool = True,
+    ) -> None:
         """Initializes the problem generator.
 
-        Initially, the generator creates n cities (x_1,y_1),...,(x_n,y_n) where n, x_i and y_i are
-        sampled independently from the provided probability distributions `n`, `x` and `y`. For each
-        (unordered) pair of cities (i,j), the distance d[i,j] between them is set to:
+        Initially, the generator creates n cities (x_1,y_1),...,(x_n,y_n) where n,
+        x_i and y_i are sampled independently from the provided probability
+        distributions `n`, `x` and `y`. For each (unordered) pair of cities (i,j),
+        the distance d[i,j] between them is set to:
 
             d[i,j] = gamma[i,j] \sqrt{(x_i - x_j)^2 + (y_i - y_j)^2}
 
         where gamma is sampled from the provided probability distribution `gamma`.
 
-        If fix_cities=True, the list of cities is kept the same for all generated instances. The
-        gamma values, and therefore also the distances, are still different.
+        If fix_cities=True, the list of cities is kept the same for all generated
+        instances. The gamma values, and therefore also the distances, are still
+        different.
 
-        By default, all distances d[i,j] are rounded to the nearest integer.  If `round=False`
-        is provided, this rounding will be disabled.
+        By default, all distances d[i,j] are rounded to the nearest integer.  If
+        `round=False` is provided, this rounding will be disabled.
 
         Arguments
         ---------
@@ -75,8 +146,8 @@ class TravelingSalesmanGenerator:
         n: rv_discrete
             Probability distribution for the number of cities.
         fix_cities: bool
-            If False, cities will be resampled for every generated instance. Otherwise, list of
-            cities will be computed once, during the constructor.
+            If False, cities will be resampled for every generated instance. Otherwise, list
+            of cities will be computed once, during the constructor.
         round: bool
             If True, distances are rounded to the nearest integer.
         """
@@ -94,14 +165,17 @@ class TravelingSalesmanGenerator:
         self.round = round
 
         if fix_cities:
+            self.fixed_n: Optional[int]
+            self.fixed_cities: Optional[np.ndarray]
             self.fixed_n, self.fixed_cities = self._generate_cities()
         else:
             self.fixed_n = None
             self.fixed_cities = None
 
-    def generate(self, n_samples):
-        def _sample():
+    def generate(self, n_samples: int) -> List[TravelingSalesmanInstance]:
+        def _sample() -> TravelingSalesmanInstance:
             if self.fixed_cities is not None:
+                assert self.fixed_n is not None
                 n, cities = self.fixed_n, self.fixed_cities
             else:
                 n, cities = self._generate_cities()
@@ -113,75 +187,7 @@ class TravelingSalesmanGenerator:
 
         return [_sample() for _ in range(n_samples)]
 
-    def _generate_cities(self):
+    def _generate_cities(self) -> Tuple[int, np.ndarray]:
         n = self.n.rvs()
         cities = np.array([(self.x.rvs(), self.y.rvs()) for _ in range(n)])
         return n, cities
-
-
-class TravelingSalesmanInstance(Instance):
-    """An instance ot the Traveling Salesman Problem.
-
-    Given a list of cities and the distance between each pair of cities, the problem asks for the
-    shortest route starting at the first city, visiting each other city exactly once, then
-    returning to the first city. This problem is a generalization of the Hamiltonian path problem,
-    one of Karp's 21 NP-complete problems.
-    """
-
-    def __init__(self, n_cities, distances):
-        super().__init__()
-        assert isinstance(distances, np.ndarray)
-        assert distances.shape == (n_cities, n_cities)
-        self.n_cities = n_cities
-        self.distances = distances
-        self.edges = [
-            (i, j) for i in range(self.n_cities) for j in range(i + 1, self.n_cities)
-        ]
-        self.varname_to_index = {f"x[{e}]": e for e in self.edges}
-
-    @overrides
-    def to_model(self):
-        model = pe.ConcreteModel()
-        model.x = pe.Var(self.edges, domain=pe.Binary)
-        model.obj = pe.Objective(
-            expr=sum(model.x[i, j] * self.distances[i, j] for (i, j) in self.edges),
-            sense=pe.minimize,
-        )
-        model.eq_degree = pe.ConstraintList()
-        model.eq_subtour = pe.ConstraintList()
-        for i in range(self.n_cities):
-            model.eq_degree.add(
-                sum(
-                    model.x[min(i, j), max(i, j)]
-                    for j in range(self.n_cities)
-                    if i != j
-                )
-                == 2
-            )
-        return model
-
-    @overrides
-    def get_variable_category(self, var_name):
-        return self.varname_to_index[var_name]
-
-    @overrides
-    def find_violated_lazy_constraints(self, model):
-        selected_edges = [e for e in self.edges if model.x[e].value > 0.5]
-        graph = nx.Graph()
-        graph.add_edges_from(selected_edges)
-        components = [frozenset(c) for c in list(nx.connected_components(graph))]
-        violations = []
-        for c in components:
-            if len(c) < self.n_cities:
-                violations += [c]
-        return violations
-
-    @overrides
-    def build_lazy_constraint(self, model, component):
-        cut_edges = [
-            e
-            for e in self.edges
-            if (e[0] in component and e[1] not in component)
-            or (e[0] not in component and e[1] in component)
-        ]
-        return model.eq_subtour.add(sum(model.x[e] for e in cut_edges) >= 2)