diff --git a/src/lib.rs b/src/lib.rs
index ca9fa2c..367344d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -64,7 +64,7 @@ fn calculate_densities_and_gaps(dataset: &[Point], factor: f32, min_cluster_size
 
     // Define thresholds for clustering and gap identification based on the mean distance and factor.
     let cluster_threshold = mean_distance / factor;
-    let gap_threshold = factor * mean_distance * 2.0;
+    let gap_threshold = factor * mean_distance;
 
     let mut results: Vec<ClusterGapInfo> = Vec::new(); // Stores the resulting clusters and gaps.
     let mut current_cluster: Vec<Point> = Vec::new(); // Temporary storage for points in the current cluster.
@@ -134,12 +134,17 @@ fn calculate_densities_and_gaps(dataset: &[Point], factor: f32, min_cluster_size
 ///
 #[pyfunction]
 fn lyagushka(_py: Python, int_list: &PyList, factor: f32, min_cluster_size: usize) -> PyResult<String> {
+    
     // Extract integers from a Python list and create a vector of Point structs.
-    let dataset: Vec<Point> = int_list.extract::<Vec<u32>>()?
+    let mut dataset: Vec<Point> = int_list.extract::<Vec<u32>>()?
                                       .into_iter()
                                       .map(Point::new)
                                       .collect();
 
+    
+    // Sort the vector
+    dataset.sort_by_key(|p| p.value);
+
     // Calculate clusters and gaps from the dataset using predefined criteria.
     let mut cluster_gap_infos = calculate_densities_and_gaps(&dataset, factor, min_cluster_size);
 
diff --git a/test.py b/test.py
index e2fd2ef..0a84dc7 100644
--- a/test.py
+++ b/test.py
@@ -35,40 +35,32 @@ dataset.sort()
 
 # calculate the anomalies in the data
 analysis_results = json.loads(lyagushka(dataset, 3.0, 7))
-print(analysis_results)
-print(len(analysis_results))
 
 # Initialize plot
 plt.figure(figsize=(10, 6))
 
-# Plot dataset points
-for point in dataset:
-    plt.plot(point, 0, 'ko')  # Plot dataset as black dots at y=0
+# Color palette for clusters and gaps
+colors = plt.cm.jet(np.linspace(0, 1, len(analysis_results)))
 
-# Process each cluster/gap for plotting
-for result in analysis_results:
+# Plot dataset points and assign colors based on cluster membership
+for i, result in enumerate(analysis_results):
     if result['num_elements'] > 0:  # It's a cluster
-        color = 'blue'  # Color for clusters
-    else:  # It's a gap
-        color = 'green'  # Color for gaps
+        points_in_cluster = [point for point in dataset if 
+                             result['centroid'] - result['span_length'] / 2 <= point <= 
+                             result['centroid'] + result['span_length'] / 2]
+        for point in points_in_cluster:
+            plt.plot(point, 0, 'o', color=colors[i])  # Plot points in cluster with the same color
 
-    # Generate start and end points for the segment
+    # Plot a line segment for the cluster/gap Z-score in the same color
     start = result['centroid'] - result['span_length'] / 2
     end = result['centroid'] + result['span_length'] / 2
     z_score = result['z_score'] if result['z_score'] is not None else 0
-
-    # Plot a line segment for the cluster/gap
-    plt.plot([start, end], [z_score, z_score], color=color, linewidth=2)
+    plt.plot([start, end], [z_score, z_score], color=colors[i], linewidth=2)
 
 # Enhancements for visualization
 plt.xlabel('Integer Value')
 plt.ylabel('Z-Score')
-plt.title('Cluster and Gap Analysis with Distinct Z-Score Curves')
+plt.title('Cluster and Gap Analysis')
 plt.grid(True)
 
-# Custom legend
-plt.plot([], [], color='blue', label='Clusters')
-plt.plot([], [], color='green', label='Gaps')
-plt.legend()
-
 plt.show()
\ No newline at end of file