diff --git a/src/main/java/sk/ableneo/openslava2023/codefix/BalloonChallenge.java b/src/main/java/sk/ableneo/openslava2023/codefix/BalloonChallenge.java
index 9b3e0a3..4b385ac 100644
--- a/src/main/java/sk/ableneo/openslava2023/codefix/BalloonChallenge.java
+++ b/src/main/java/sk/ableneo/openslava2023/codefix/BalloonChallenge.java
@@ -6,26 +6,24 @@
 
 import java.util.ArrayList;
 import java.util.LinkedList;
+import java.util.List;
 
 // TODO: fix the class
 public class BalloonChallenge {
 
     public LinkedList findOrdersForProduct(Product p, boolean debug) {
-        ArrayList l = new ArrayList();
-        ArrayList list = getAllOrders();
-        for (int i = 0; i < list.size(); i++) {
-            Order order = (Order) list.get(i);
-            boolean found = false;
-            if (order.getProducts().size() > 0) {
-                for (int j = 0; j <= order.getProducts().size(); j++) {
-                    Product p2 = (Product) order.getProducts().get(j);
-                    found = (p2 == p);
-                }
-                if (found && order != null)
-                    l.add(order);
-            }
+        ArrayList results = new ArrayList();
+        ArrayList allOrders = getAllOrders();
+
+        for (int i = 0; i < allOrders.size(); i++) {
+            Order order = (Order) allOrders.get(i);
+            boolean found = order.getProducts().contains(p);
+
+            if (found)
+                results.add(order);
         }
-        return new LinkedList(l);
+        
+        return new LinkedList(results);
     }
 
     public ArrayList getAllOrders() {