1 files changed, 30 insertions, 0 deletions
diff --git a/src/retain.rs b/src/retain.rs
new file mode 100644
index 0000000..bc3934a
--- /dev/null
+++ b/src/retain.rs
@@ -0,0 +1,30 @@
+use hashbrown::HashSet;
+use num_complex::Complex;
+
+#[inline(always)]
+pub fn retain_top_n_magnitudes(fft_real_signal: &mut [Complex<f32>], n: usize) {
+    // you'd be keeping the entire signal
+    if n >= fft_real_signal.len() {
+        return;
+    }
+
+    let mut indexed = fft_real_signal
+        .iter()
+        .copied()
+        .enumerate()
+        .collect::<Vec<(usize, Complex<f32>)>>();
+
+    let target = fft_real_signal.len() - n;
+    indexed.select_nth_unstable_by(target, |(_, c0), (_, c1)| c0.norm().total_cmp(&c1.norm()));
+
+    let top_indices = indexed[target..]
+        .iter()
+        .map(|&(i, _)| i)
+        .collect::<HashSet<usize>>();
+
+    for (i, val) in fft_real_signal.iter_mut().enumerate() {
+        if !top_indices.contains(&i) {
+            *val = Complex::ZERO;
+        }
+    }
+}