diff --git a/Data Generation and Likelihood Fit.ipynb b/Data Generation and Likelihood Fit.ipynb
index 22e3804..b65c144 100644
--- a/Data Generation and Likelihood Fit.ipynb	
+++ b/Data Generation and Likelihood Fit.ipynb	
@@ -32,6 +32,7 @@
     "import random as rnd\n",
     "import scipy.stats as stats\n",
     "import scipy.optimize as opt\n",
+    "import scipy.sparse\n",
     "import json as json\n",
     "import matplotlib as mpl\n",
     "from math import exp\n",
@@ -526,17 +527,8 @@
     "    # Initialization of the state-transition matrices: describe the state-transition probabilities\n",
     "    # if the maintenance cost is incurred, and regenerate the state to 0 if the replacement cost\n",
     "    # is incurred.\n",
-    "    ST_mat = np.zeros((S, S))\n",
-    "    p = np.array(p) \n",
-    "    for i in range(S):\n",
-    "        for j, _p in enumerate(p):\n",
-    "            if i + j < S-1:\n",
-    "                ST_mat[i+j][i] = _p\n",
-    "                \n",
-    "            elif i + j == S-1:\n",
-    "                ST_mat[S-1][i] = p[j:].sum()\n",
-    "            else:\n",
-    "                pass\n",
+    "    ST_mat=scipy.sparse.diags(p, range(len(p)), shape=(S, S)).toarray()\n",
+    "    ST_mat[-2:,-1] = [1 - p[1], 1]\n",
     "    \n",
     "    R_mat = np.vstack((np.ones((1, S)),np.zeros((S-1, S))))        \n",
     "            \n",
@@ -1211,15 +1203,8 @@
     "        # A (SxS) matrix indicating the probability of a bus transitioning\n",
     "        # from a state s to a state s' (used to compute maintenance utility)\n",
     "        \n",
-    "        self.trans_mat = np.zeros((S, S))\n",
-    "        for i in range(S):\n",
-    "            for j, _p in enumerate(self.p):\n",
-    "                if i + j < S-1:\n",
-    "                    self.trans_mat[i+j][i] = _p\n",
-    "                elif i + j == S-1:\n",
-    "                    self.trans_mat[S-1][i] = p[j:].sum()\n",
-    "                else:\n",
-    "                    pass\n",
+    "        self.trans_mat=scipy.sparse.diags(p, range(len(p)), shape=(S, S)).toarray()\n",
+    "        self.trans_mat[-2:,-1] = [1 - p[1], 1]\n",
     "\n",
     "        # A second (SxS) matrix which regenerates the bus' state to 0 with\n",
     "        # certainty (used to compute the replacement utility)\n",