From b0d8050b110d4623c102c8107994551d6f98bbaf Mon Sep 17 00:00:00 2001
From: Palash Tyagi <23239946+Magnus167@users.noreply.github.com>
Date: Sat, 26 Jul 2025 13:26:44 +0100
Subject: [PATCH] add test as examples

---
 examples/k_means.rs | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/examples/k_means.rs b/examples/k_means.rs
index c1144c0..4bc1144 100644
--- a/examples/k_means.rs
+++ b/examples/k_means.rs
@@ -49,3 +49,17 @@ fn customer_spend_example() {
     let pred = model.predict(&new_customers);
     println!("Cluster of new customers: {:?}", pred);
 }
+
+#[test]
+fn k_means_store_locations() {
+    let raw = vec![
+        1.0, 2.0, 1.5, 1.8, 5.0, 8.0, 8.0, 8.0, 1.0, 0.6, 9.0, 11.0, 8.0, 2.0, 10.0, 2.0, 9.0, 3.0,
+    ];
+    let x = Matrix::from_rows_vec(raw, 9, 2);
+    let (model, labels) = KMeans::fit(&x, 2, 100, 1e-4);
+    assert_eq!(labels.len(), 9);
+    assert_eq!(model.centroids.rows(), 2);
+    let new_points = Matrix::from_rows_vec(vec![0.0, 0.0, 8.0, 3.0], 2, 2);
+    let pred = model.predict(&new_points);
+    assert_eq!(pred.len(), 2);
+}