diff --git a/src/decompression.jl b/src/decompression.jl
index c39342b6..dfc29575 100644
--- a/src/decompression.jl
+++ b/src/decompression.jl
@@ -592,6 +592,8 @@ function decompress!(
         diagonal_nzind,
         lower_triangle_offsets,
         upper_triangle_offsets,
+        compressed_diag,
+        compressed_indices,
         buffer,
     ) = result
     (; S) = ag
@@ -608,21 +610,24 @@ function decompress!(
     # Recover the diagonal coefficients of A
     if has_diagonal(ag)
         if uplo == :L
-            for i in diagonal_indices
+            for (k, i) in enumerate(diagonal_indices)
                 # A[i, i] is the first element in column i
                 nzind = A_colptr[i]
-                nzA[nzind] = B[i, color[i]]
+                pos = compressed_diag[k]
+                nzA[nzind] = B[pos]
             end
         elseif uplo == :U
-            for i in diagonal_indices
+            for (k, i) in enumerate(diagonal_indices)
                 # A[i, i] is the last element in column i
                 nzind = A_colptr[i + 1] - 1
-                nzA[nzind] = B[i, color[i]]
+                pos = compressed_diag[k]
+                nzA[nzind] = B[pos]
             end
         else  # uplo == :F
             for (k, i) in enumerate(diagonal_indices)
                 nzind = diagonal_nzind[k]
-                nzA[nzind] = B[i, color[i]]
+                pos = compressed_diag[k]
+                nzA[nzind] = B[pos]
             end
         end
     end
@@ -642,7 +647,8 @@ function decompress!(
 
         for (i, j) in reverse_bfs_orders[k]
             counter += 1
-            val = B[i, color[j]] - buffer_right_type[i]
+            pos = compressed_indices[counter]
+            val = B[pos] - buffer_right_type[i]
             buffer_right_type[j] = buffer_right_type[j] + val
 
             #! format: off
diff --git a/src/result.jl b/src/result.jl
index 8bbefddf..96bb7541 100644
--- a/src/result.jl
+++ b/src/result.jl
@@ -283,6 +283,8 @@ struct TreeSetColoringResult{M<:AbstractMatrix,G<:AdjacencyGraph,V,R} <:
     diagonal_nzind::Vector{Int}
     lower_triangle_offsets::Vector{Int}
     upper_triangle_offsets::Vector{Int}
+    compressed_diag::Vector{Int}
+    compressed_indices::Vector{Int}
     buffer::Vector{R}
 end
 
@@ -302,6 +304,7 @@ function TreeSetColoringResult(
     # Vector for the decompression of the diagonal coefficients
     diagonal_indices = Int[]
     diagonal_nzind = Int[]
+    compressed_diag = Int[]
     ndiag = 0
 
     if has_diagonal(ag)
@@ -311,6 +314,7 @@ function TreeSetColoringResult(
                 if i == j
                     push!(diagonal_indices, i)
                     push!(diagonal_nzind, k)
+                    push!(compressed_diag, (color[i] - 1) * nvertices + i)
                     ndiag += 1
                 end
             end
@@ -321,8 +325,9 @@ function TreeSetColoringResult(
     nedges = (nnz(S) - ndiag) ÷ 2
     lower_triangle_offsets = Vector{Int}(undef, nedges)
     upper_triangle_offsets = Vector{Int}(undef, nedges)
+    compressed_indices = Vector{Int}(undef, nedges)
 
-    # Index in lower_triangle_offsets and upper_triangle_offsets
+    # Index in compressed_indices, lower_triangle_offsets and upper_triangle_offsets
     index_offsets = 0
 
     for k in eachindex(reverse_bfs_orders)
@@ -334,6 +339,8 @@ function TreeSetColoringResult(
             col_j = view(rv, nzrange(S, j))
             index_offsets += 1
 
+            compressed_indices[index_offsets] = (color[j] - 1) * nvertices + i
+
             #! format: off
             # S[i,j] is in the lower triangular part of S
             if in_triangle(i, j, :L)
@@ -373,6 +380,8 @@ function TreeSetColoringResult(
         diagonal_nzind,
         lower_triangle_offsets,
         upper_triangle_offsets,
+        compressed_diag,
+        compressed_indices,
         buffer,
     )
 end