From d271178165a7351d902478e6a544a1287e58e007 Mon Sep 17 00:00:00 2001
From: Andrei Ovcharenko <a@krot.name>
Date: Fri, 31 Jul 2026 17:44:30 +0300
Subject: [PATCH] Validate DorogayaSet input bounds

---
 .../java/algorithms/sprint6/DorogayaSet.java  | 51 ++++++++++++++++---
 .../java/algorithms/AlgorithmCliTest.java     |  6 +++
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/src/main/java/algorithms/sprint6/DorogayaSet.java b/src/main/java/algorithms/sprint6/DorogayaSet.java
index a1a7c20..d47b1fb 100644
--- a/src/main/java/algorithms/sprint6/DorogayaSet.java
+++ b/src/main/java/algorithms/sprint6/DorogayaSet.java
@@ -43,8 +43,14 @@
 public class DorogayaSet {
 
     static final String FAIL = "Oops! I did it again";
+    private static final int MAX_VERTICES = 200_000;
+    private static final int MAX_EDGES = 200_000;
 
     static long solve(int n, Edge[] edges) {
+        if (!isValidGraph(n, edges)) {
+            return -1;
+        }
+
         Arrays.sort(edges, (a, b) -> Integer.compare(b.weight, a.weight));
 
         DSU dsu = new DSU(n);
@@ -69,6 +75,24 @@ static long solve(int n, Edge[] edges) {
         return totalWeight;
     }
 
+    private static boolean isValidGraph(int n, Edge[] edges) {
+        if (n < 1 || n > MAX_VERTICES || edges == null || edges.length > MAX_EDGES) {
+            return false;
+        }
+
+        for (Edge edge : edges) {
+            if (edge == null || !isValidVertex(edge.from, n) || !isValidVertex(edge.to, n)) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    private static boolean isValidVertex(int vertex, int n) {
+        return vertex >= 1 && vertex <= n;
+    }
+
     static final class Edge {
         final int from;
         final int to;
@@ -236,17 +260,28 @@ private static void run() throws Exception {
         int n = in.nextInt();
         int m = in.nextInt();
 
-        Edge[] edges = new Edge[m];
+        long answer = -1;
 
-        for (int i = 0; i < m; i++) {
-            int from = in.nextInt();
-            int to = in.nextInt();
-            int weight = in.nextInt();
+        if (n >= 1 && n <= MAX_VERTICES && m >= 0 && m <= MAX_EDGES) {
+            Edge[] edges = new Edge[m];
+            boolean validEdges = true;
 
-            edges[i] = new Edge(from, to, weight);
-        }
+            for (int i = 0; i < m; i++) {
+                int from = in.nextInt();
+                int to = in.nextInt();
+                int weight = in.nextInt();
+
+                if (!isValidVertex(from, n) || !isValidVertex(to, n)) {
+                    validEdges = false;
+                }
+
+                edges[i] = new Edge(from, to, weight);
+            }
 
-        long answer = solve(n, edges);
+            if (validEdges) {
+                answer = solve(n, edges);
+            }
+        }
 
         if (answer == -1) {
             out.writeString(FAIL);
diff --git a/src/test/java/algorithms/AlgorithmCliTest.java b/src/test/java/algorithms/AlgorithmCliTest.java
index 0da7388..80aede9 100644
--- a/src/test/java/algorithms/AlgorithmCliTest.java
+++ b/src/test/java/algorithms/AlgorithmCliTest.java
@@ -120,6 +120,12 @@ private static Stream<Arguments> stdioCases() {
                 arguments("algorithms.sprint6.DorogayaSet", "run",
                         "4 4%n1 2 5%n1 3 6%n2 4 8%n3 4 3%n".formatted(),
                         "19%n".formatted()),
+                arguments("algorithms.sprint6.DorogayaSet", "run",
+                        "5 -1%n".formatted(),
+                        "Oops! I did it again%n".formatted()),
+                arguments("algorithms.sprint6.DorogayaSet", "run",
+                        "2 1%n3 1 42%n".formatted(),
+                        "Oops! I did it again%n".formatted()),
                 arguments("algorithms.sprint6.WaterWorld", "run",
                         "3 3%n#.#%n.#.%n#.#%n".formatted(),
                         "5 1%n".formatted()),