diff --git a/src/DynamicModel.cc b/src/DynamicModel.cc
index 1525b30ab301ae83b2f97e1a0168f06b0b2b32fd..20bbde1ea934e722a4818043a5c17bbf36bd4cfb 100644
--- a/src/DynamicModel.cc
+++ b/src/DynamicModel.cc
@@ -89,7 +89,6 @@ DynamicModel::DynamicModel(const DynamicModel &m) :
   block_det_exo_index{m.block_det_exo_index},
   block_other_endo_index{m.block_other_endo_index},
   variable_block_lead_lag{m.variable_block_lead_lag},
-  equation_block{m.equation_block},
   var_expectation_functions_to_write{m.var_expectation_functions_to_write}
 {
   copyHelper(m);
@@ -143,7 +142,6 @@ DynamicModel::operator=(const DynamicModel &m)
   block_det_exo_index = m.block_det_exo_index;
   block_other_endo_index = m.block_other_endo_index;
   variable_block_lead_lag = m.variable_block_lead_lag;
-  equation_block = m.equation_block;
   var_expectation_functions_to_write = m.var_expectation_functions_to_write;
 
   copyHelper(m);
@@ -3201,13 +3199,7 @@ DynamicModel::writeOutput(ostream &output, const string &basename, bool block_de
           for (int other_endo : other_endogenous)
             if (other_endo >= 0)
               {
-                bool OK = true;
-                int j;
-                for (j = 0; j < block && OK; j++)
-                  for (int k = 0; k < blocks[j].size && OK; k++)
-                    OK = other_endo != getBlockVariableID(j, k);
-                if (!OK)
-                  output << " " << j;
+                output << " " << endo2block[other_endo]+1;
                 i++;
               }
           output << "];" << endl;
@@ -4838,14 +4830,12 @@ DynamicModel::computingPass(bool jacobianExo, int derivsOrder, int paramsDerivsO
       if (!no_tmp_terms)
         computeTemporaryTermsOrdered();
       int k = 0;
-      equation_block.resize(equations.size());
       variable_block_lead_lag.clear();
       variable_block_lead_lag.resize(equations.size());
       for (int i = 0; i < static_cast<int>(blocks.size()); i++)
         {
           for (int j = 0; j < blocks[i].size; j++)
             {
-              equation_block[eq_idx_block2orig[k]] = i;
               int l = endo_idx_block2orig[k];
               variable_block_lead_lag[l] = { i, variable_lag_lead[l].first, variable_lag_lead[l].second };
               k++;
@@ -5047,15 +5037,7 @@ DynamicModel::computeChainRuleJacobian()
 void
 DynamicModel::collect_block_first_order_derivatives()
 {
-  //! vector for an equation or a variable indicates the block number
-  vector<int> equation_2_block(eq_idx_block2orig.size()), variable_2_block(endo_idx_block2orig.size());
   int nb_blocks = blocks.size();
-  for (int block = 0; block < nb_blocks; block++)
-    for (int i = 0; i < blocks[block].size; i++)
-      {
-        equation_2_block[getBlockEquationID(block, i)] = block;
-        variable_2_block[getBlockVariableID(block, i)] = block;
-      }
   other_endo_block = vector<lag_var_t>(nb_blocks);
   exo_block = vector<lag_var_t>(nb_blocks);
   exo_det_block = vector<lag_var_t>(nb_blocks);
@@ -5072,14 +5054,14 @@ DynamicModel::collect_block_first_order_derivatives()
       int eq = indices[0];
       int var = symbol_table.getTypeSpecificID(getSymbIDByDerivID(indices[1]));
       int lag = getLagByDerivID(indices[1]);
-      int block_eq = equation_2_block[eq];
+      int block_eq = eq2block[eq];
       int block_var = 0;
       derivative_t tmp_derivative;
       lag_var_t lag_var;
       switch (getTypeByDerivID(indices[1]))
         {
         case SymbolType::endogenous:
-          block_var = variable_2_block[var];
+          block_var = endo2block[var];
           if (block_eq == block_var)
             {
               if (lag < 0 && lag < -endo_max_leadlag_block[block_eq].first)
diff --git a/src/ModelTree.cc b/src/ModelTree.cc
index f92e4db9ca6153caad1b696deae32df94b37adad..3b45efa28dc162494a29e3b0f96c82af1faf8fb8 100644
--- a/src/ModelTree.cc
+++ b/src/ModelTree.cc
@@ -162,6 +162,8 @@ ModelTree::ModelTree(const ModelTree &m) :
   exo_det_max_leadlag_block{m.exo_det_max_leadlag_block},
   max_leadlag_block{m.max_leadlag_block},
   blocks{m.blocks},
+  endo2block{m.endo2block},
+  eq2block{m.eq2block},
   is_equation_linear{m.is_equation_linear},
   endo2eq{m.endo2eq},
   epilogue{m.epilogue},
@@ -214,6 +216,8 @@ ModelTree::operator=(const ModelTree &m)
   exo_det_max_leadlag_block = m.exo_det_max_leadlag_block;
   max_leadlag_block = m.max_leadlag_block;
   blocks = m.blocks;
+  endo2block = m.endo2block;
+  eq2block = m.eq2block;
   is_equation_linear = m.is_equation_linear;
   endo2eq = m.endo2eq;
   epilogue = m.epilogue;
@@ -412,7 +416,8 @@ ModelTree::select_non_linear_equations_and_variables()
   prologue = 0;
   epilogue = 0;
 
-  vector<int> endo2block(endo2eq.size(), 1); // The 1 is a dummy value, distinct from 0
+  endo2block.resize(endo2eq.size(), 1); // The 1 is a dummy value, distinct from 0
+  eq2block.resize(endo2eq.size(), 1);
   int i = 0;
   for (int endo = 0; endo < static_cast<int>(endo2eq.size()); endo++)
     {
@@ -421,7 +426,8 @@ ModelTree::select_non_linear_equations_and_variables()
         {
           eq_idx_block2orig[i] = eq;
           endo_idx_block2orig[i] = endo;
-          endo2block[i] = 0;
+          endo2block[endo] = 0;
+          eq2block[eq] = 0;
           i++;
         }
     }
@@ -432,7 +438,7 @@ ModelTree::select_non_linear_equations_and_variables()
   blocks[0].size = i;
   blocks[0].mfs_size = i;
 
-  auto [equation_lag_lead, variable_lag_lead] = getVariableLeadLagByBlock(endo2block);
+  auto [equation_lag_lead, variable_lag_lead] = getVariableLeadLagByBlock();
 
   for (int i = 0; i < blocks[0].size; i++)
     {
@@ -620,28 +626,17 @@ ModelTree::equationTypeDetermination(const map<tuple<int, int, int>, expr_t> &fi
 }
 
 pair<lag_lead_vector_t, lag_lead_vector_t>
-ModelTree::getVariableLeadLagByBlock(const vector<int> &endo2simblock) const
+ModelTree::getVariableLeadLagByBlock() const
 {
   int nb_endo = symbol_table.endo_nbr();
 
-  auto belong_to_same_block = [&](int endo, int eq)
-                              {
-                                int endo2 = endo_idx_orig2block[endo];
-                                int eq2 = eq_idx_orig2block[eq];
-                                if (endo2 < prologue || endo2 >= nb_endo-epilogue
-                                    || eq2 < prologue || eq2 >= nb_endo-epilogue)
-                                  return endo2 == eq2;
-                                else
-                                  return endo2simblock[endo2-prologue] == endo2simblock[eq2-prologue];
-                              };
-
   lag_lead_vector_t variable_lag_lead(nb_endo, { 0, 0 }), equation_lag_lead(nb_endo, { 0, 0 });
   for (int eq = 0; eq < nb_endo; eq++)
     {
       set<pair<int, int>> endos_and_lags;
       equations[eq]->collectEndogenous(endos_and_lags);
       for (auto [endo, lag] : endos_and_lags)
-        if (belong_to_same_block(endo, eq))
+        if (endo2block[endo] == eq2block[eq])
           {
             variable_lag_lead[endo].first = max(variable_lag_lead[endo].first, -lag);
             variable_lag_lead[endo].second = max(variable_lag_lead[endo].second, lag);
@@ -684,40 +679,45 @@ ModelTree::computeBlockDecompositionAndFeedbackVariablesForEachBlock()
 
   blocks.clear();
   blocks.resize(num_blocks);
+  endo2block.resize(nb_var);
+  eq2block.resize(nb_var);
 
-  // Initialize size and mfs_size for prologue and epilogue
+  // Initialize size and mfs_size for prologue and epilogue, plus eq/endo→block mappings
   for (int i = 0; i < prologue; i++)
     {
       blocks[i].size = 1;
       blocks[i].mfs_size = 1;
+      endo2block[endo_idx_block2orig[i]] = i;
+      eq2block[eq_idx_block2orig[i]] = i;
     }
   for (int i = 0; i < epilogue; i++)
     {
-      blocks[prologue+num_simblocks+i].size = 1;
-      blocks[prologue+num_simblocks+i].mfs_size = 1;
+      int blk = prologue+num_simblocks+i;
+      int var_eq = prologue+nb_simvars+i;
+      blocks[blk].size = 1;
+      blocks[blk].mfs_size = 1;
+      endo2block[endo_idx_block2orig[var_eq]] = blk;
+      eq2block[eq_idx_block2orig[var_eq]] = blk;
     }
 
   // Compute size and list of equations for simultaneous blocks
   vector<set<int>> eqs_in_simblock(num_simblocks);
   for (int i = 0; i < static_cast<int>(endo2simblock.size()); i++)
     {
-      blocks[prologue+endo2simblock[i]].size++;
       eqs_in_simblock[endo2simblock[i]].insert(i);
+      int blk = prologue+endo2simblock[i];
+      blocks[blk].size++;
+      endo2block[endo_idx_block2orig[prologue+i]] = blk;
+      eq2block[eq_idx_block2orig[prologue+i]] = blk;
     }
 
   // Determine the dynamic structure of each block
-  auto [equation_lag_lead, variable_lag_lead] = getVariableLeadLagByBlock(endo2simblock);
+  auto [equation_lag_lead, variable_lag_lead] = getVariableLeadLagByBlock();
 
   for (int var = 0; var < nb_var; var++)
     {
       auto [max_lag, max_lead] = variable_lag_lead[endo_idx_block2orig[var]];
-      int blk;
-      if (var < prologue)
-        blk = var;
-      else if (var >= nb_var - epilogue)
-        blk = var - nb_simvars + num_simblocks;
-      else
-        blk = endo2simblock[var-prologue] + prologue;
+      int blk = endo2block[endo_idx_block2orig[var]];
 
       if (max_lag != 0 && max_lead != 0)
         blocks[blk].n_mixed++;
@@ -829,10 +829,7 @@ ModelTree::reduceBlocksAndTypeDetermination(bool linear_decomposition)
           set<pair<int, int>> endos_and_lags;
           equations[eq_idx_block2orig[eq]]->collectEndogenous(endos_and_lags);
           for (const auto &[endo, lag] : endos_and_lags)
-            if (linear_decomposition ||
-                find(endo_idx_block2orig.begin()+first_eq,
-                        endo_idx_block2orig.begin()+first_eq+blocks[i].size, endo)
-                != endo_idx_block2orig.begin()+first_eq+blocks[i].size)
+            if (linear_decomposition || endo2block[endo] == i)
               {
                 max_lead = max(lag, max_lead);
                 max_lag = max(-lag, max_lag);
@@ -913,6 +910,12 @@ ModelTree::reduceBlocksAndTypeDetermination(bool linear_decomposition)
                   blocks[i-1].n_backward += blocks[i].n_backward;
                   blocks[i-1].n_mixed += blocks[i].n_mixed;
                   blocks.erase(blocks.begin()+i);
+                  for (auto &b : endo2block)
+                    if (b >= i)
+                      b--;
+                  for (auto &b : eq2block)
+                    if (b >= i)
+                      b--;
                   i--;
                   continue;
                 }
diff --git a/src/ModelTree.hh b/src/ModelTree.hh
index 2bffcf811748304040b9e27f1c97040c6aa80234..d7a2adb91e1941b3b6a64654588435b563f128dd 100644
--- a/src/ModelTree.hh
+++ b/src/ModelTree.hh
@@ -200,6 +200,12 @@ protected:
   // Stores various informations on the blocks
   vector<BlockInfo> blocks;
 
+  // Maps endogenous type-specific IDs to the block number to which it belongs
+  vector<int> endo2block;
+  /* Maps (original) equation number to the block number to which it belongs.
+     It verifies: ∀i, eq2block[endo2eq[i]] = endo2block[i] */
+  vector<int> eq2block;
+
   //! the file containing the model and the derivatives code
   ofstream code_file;
 
@@ -296,7 +302,8 @@ protected:
   /* Compute the block decomposition and for a non-recusive block find the minimum feedback set
 
      Initializes the “blocks” structure, and fills the following fields: size,
-     mfs_size, n_static, n_forward, n_backward, n_mixed. */
+     mfs_size, n_static, n_forward, n_backward, n_mixed.
+     Also initializes the endo2block and eq2block structures. */
   lag_lead_vector_t computeBlockDecompositionAndFeedbackVariablesForEachBlock();
   /* Reduce the number of block by merging the same type of equations in the
      prologue and the epilogue, and determine the type of each block.
@@ -312,7 +319,7 @@ protected:
      The 2nd output gives, for each type-specific endo IDs, its (max_lag,
      max_lead) across all its occurences inside the equations of the block to
      which it belongs. */
-  pair<lag_lead_vector_t, lag_lead_vector_t> getVariableLeadLagByBlock(const vector<int> &endo2simblock) const;
+  pair<lag_lead_vector_t, lag_lead_vector_t> getVariableLeadLagByBlock() const;
   //! For each equation determine if it is linear or not
   void equationLinear(const map<tuple<int, int, int>, expr_t> &first_order_endo_derivatives);
   //! Print an abstract of the block structure of the model