From 67133d0e7cfeadb14d53ba5c0f8966f0d6b14425 Mon Sep 17 00:00:00 2001
From: Dylan Asmar <asmar@stanford.edu>
Date: Mon, 22 Jul 2024 14:31:49 -0700
Subject: [PATCH] Updated `prune_strictly_dominated!` to reduce allocations

---
 src/prune.jl | 41 +++++++++++++++++++++++++++++++----------
 1 file changed, 31 insertions(+), 10 deletions(-)

diff --git a/src/prune.jl b/src/prune.jl
index 58b664b..395da7d 100644
--- a/src/prune.jl
+++ b/src/prune.jl
@@ -51,10 +51,10 @@ end
 
 function intersection_distance(α1, α2, b)
     dot_sum = 0.0
-    I,B = b.nzind, b.nzval
+    I, B = b.nzind, b.nzval
     for _i ∈ eachindex(I)
         i = I[_i]
-        dot_sum += (α1[i] - α2[i])*B[_i]
+        dot_sum += (α1[i] - α2[i]) * B[_i]
     end
     s = 0.0
     for i ∈ eachindex(α1, α2)
@@ -122,26 +122,47 @@ end
 
 function prune_strictly_dominated!(tree::SARSOPTree, eps=1e-10)
     Γ = tree.Γ
-    Γ_new_idxs = Int[]
+    Γ_new_idxs = Vector{Int}(undef, length(Γ))
+    keep = trues(length(Γ))
 
+    idx_count = 0
     for (α_try_idx, α_try) in enumerate(Γ)
-        marked_for_deletion = falses(length(Γ_new_idxs))
         dominated = false
-        for (jj, α_in_idx) in enumerate(Γ_new_idxs)
+        for jj in 1:idx_count
+            α_in_idx = Γ_new_idxs[jj]
             α_in = Γ[α_in_idx]
             if strictly_dominates(α_try, α_in, eps)
-                marked_for_deletion[jj] = true
+                keep[jj] = false
             elseif strictly_dominates(α_in, α_try, eps)
                 dominated = true
                 break
             end
         end
         if !dominated
-            Γ_new_idxs = Γ_new_idxs[.!marked_for_deletion]
-            push!(Γ_new_idxs, α_try_idx)
+            new_idx_count = 0
+            for jj in 1:idx_count
+                if keep[jj]
+                    new_idx_count += 1
+                    Γ_new_idxs[new_idx_count] = Γ_new_idxs[jj]
+                end
+            end
+            new_idx_count += 1
+            Γ_new_idxs[new_idx_count] = α_try_idx
+            idx_count = new_idx_count
+            fill!(keep, true)
         end
     end
 
-    Γ_idxs_to_delete = setdiff(1:length(Γ), Γ_new_idxs)
-    deleteat!(Γ, Γ_idxs_to_delete)
+    resize!(Γ_new_idxs, idx_count)
+
+    to_delete = trues(length(Γ))
+    for idx in Γ_new_idxs
+        to_delete[idx] = false
+    end
+
+    for ii in length(Γ):-1:1
+        if to_delete[ii]
+            deleteat!(Γ, ii)
+        end
+    end
 end