diff --git a/lab5/zad2/src/main.rs b/lab5/zad2/src/main.rs
index 45654e9..4d8da34 100644
--- a/lab5/zad2/src/main.rs
+++ b/lab5/zad2/src/main.rs
@@ -22,31 +22,23 @@ impl Graph {
     pub fn new(size: usize) -> Self {
         Self {
             size,
-            edges: (0..size)
-                .map(|i|
-                    (0..size)
-                        .map(|j|
-                            if i == j { 0.0 }
-                            else { thread_rng().gen::<f64>() }
-                        )
-                        .collect::<Vec<_>>()
-                )
-                .collect::<Vec<_>>(),
-        }
-    }
+            edges: {
+                let mut edges = vec![vec![0.0; size]; size];
 
-    fn min(&self, key: &Vec<f64>, mst: &Vec<bool>) -> usize {
-        let mut min = f64::INFINITY;
-        let mut min_index = usize::MAX;
+                for i in 0..size {
+                    for j in 0..size {
+                        if i == j { continue }
 
-        for i in 0..self.size {
-            if mst[i] == false && key[i] < min {
-                min = key[i];
-                min_index = i;
+                        let w = thread_rng().gen::<f64>();
+
+                        edges[i][j] = w;
+                        edges[j][i] = w;
+                    }
+                }
+
+                edges
             }
         }
-
-        min_index
     }
 
     pub fn prim(&self) -> Vec<Edge> {
@@ -54,25 +46,38 @@ impl Graph {
 
         let mut key = vec![f64::INFINITY; self.size];
 
-        let mut mst = vec![false; self.size];
+        let mut mst_set = vec![false; self.size];
 
         key[0] = 0.0;
         parent[0] = usize::MAX;
 
         for _ in 0..(self.size - 1) {
-            let u = self.min(&key, &mst);
-            mst[u] = true;
+            let u = {
+                let mut min_index = 0;
+                let mut min = f64::MAX;
 
-            for j in 0..self.size {
-                if self.edges[u][j] != 0.0 && mst[j] == false && self.edges[u][j] < key[j] {
-                    parent[j] = u;
-                    key[j] = self.edges[u][j];
+                for k in 0..self.size {
+                    if mst_set[k] == false && key[k] < min {
+                        min = key[k];
+                        min_index = k;
+                    }
+                }
+
+                min_index
+            };
+
+            mst_set[u] = true;
+
+            for v in 0..self.size {
+                let w = self.edges[u][v];
+                if u != v && mst_set[v] == false && w < key[v] {
+                    parent[v] = u;
+                    key[v] = w;
                 }
             }
         }
 
-        (0..self.size)
-            .filter(|i| parent[*i] != usize::MAX)
+        (1..self.size)
             .map(|i| Edge {
                 u: parent[i],
                 v: i,
@@ -101,18 +106,18 @@ impl Graph {
         }
 
         if let Some(max) = children_rounds.iter().max() {
-            return max + 1;
+            return *max.max(&(children_rounds.len() as u64)) + 1;
         } else {
-            return 1;
+            return 0;
         }
     }
 
-    fn calculate_rounds(&self, mst: &Vec<Edge>) -> u64 {
+    fn calculate_rounds(&self, mst: &Vec<Edge>) -> (usize, u64) {
         let mut visited = vec![false; self.size];
 
-        let rand_vertex_index = thread_rng().gen_range(0..self.size);
+        let start_vertex = thread_rng().gen_range(0..self.size);
 
-        return self.make_round(rand_vertex_index, &mut visited, mst);
+        (start_vertex, self.make_round(start_vertex, &mut visited, mst))
     }
 }
 
@@ -185,7 +190,7 @@ fn test() {
 
             let mst = graph.prim();
 
-            let rounds = graph.calculate_rounds(&mst);
+            let (_, rounds) = graph.calculate_rounds(&mst);
 
             rounds_clone.get_mut(&n).unwrap().push(rounds);
         });
@@ -238,7 +243,9 @@ fn show() {
 
     print_mst(&mst_prim);
 
-    println!("rounds: {}", graph.calculate_rounds(&mst_prim));
+    let (start, rounds) = graph.calculate_rounds(&mst_prim);
+
+    println!("start: {start}, rounds: {rounds}");
 }
 
 fn main() {