diff --git a/core-java/src/main/java/com/baeldung/graph/Graph.java b/core-java/src/main/java/com/baeldung/graph/Graph.java
index 7095f64709..43b5c0aa08 100644
--- a/core-java/src/main/java/com/baeldung/graph/Graph.java
+++ b/core-java/src/main/java/com/baeldung/graph/Graph.java
@@ -1,67 +1,57 @@
 package com.baeldung.graph;
 
+import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.stream.Collectors;
 
 public class Graph {
-    private Set<Vertex> vertices;
-    private Set<Edge> edges;
-    private Map<Vertex, Set<Edge>> adjVertices;
+    private Map<Vertex, List<Vertex>> adjVertices;
 
     Graph() {
-        this.vertices = new HashSet<Vertex>();
-        this.edges = new HashSet<Edge>();
-        this.adjVertices = new HashMap<Vertex, Set<Edge>>();
+        this.adjVertices = new HashMap<Vertex, List<Vertex>>();
     }
 
-    boolean addVertex(String label) {
-        return vertices.add(new Vertex(label));
+    void addVertex(String label) {
+        adjVertices.putIfAbsent(new Vertex(label), new ArrayList<>());
     }
 
-    boolean removeVertex(String label) {
-        return vertices.remove(new Vertex(label));
-    }
-
-    boolean addEdge(String label1, String label2) {
-        Vertex v1 = new Vertex(label1);
-        Vertex v2 = new Vertex(label2);
-        Edge e = new Edge(v1, v2);
-        if (this.edges.add(e)) {
-            adjVertices.putIfAbsent(v1, new HashSet<>());
-            adjVertices.putIfAbsent(v2, new HashSet<>());
-            adjVertices.get(v1)
-                .add(e);
-            adjVertices.get(v2)
-                .add(e);
-        }
-        return true;
-    }
-
-    boolean removeEdge(String label1, String label2) {
-        Vertex v1 = new Vertex(label1);
-        Vertex v2 = new Vertex(label2);
-        Edge e = new Edge(v1, v2);
-        if (this.edges.remove(e)) {
-            Set<Edge> eV1 = adjVertices.get(v1);
-            Set<Edge> eV2 = adjVertices.get(v2);
-
-            if (eV1 != null)
-                eV1.remove(e);
-            if (eV2 != null)
-                eV2.remove(e);
-        }
-        return true;
-    }
-
-    Set<Vertex> getAdjVertices(String label) {
+    void removeVertex(String label) {
         Vertex v = new Vertex(label);
-        return adjVertices.get(v)
-            .stream()
-            .map(e -> e.v1.equals(v) ? e.v2 : e.v1)
-            .collect(Collectors.toSet());
+        adjVertices.values().stream().map(e -> e.remove(v)).collect(Collectors.toList());
+        adjVertices.remove(new Vertex(label));
+    }
+
+    void addEdge(String label1, String label2) {
+        Vertex v1 = new Vertex(label1);
+        Vertex v2 = new Vertex(label2);
+        adjVertices.get(v1).add(v2);
+        adjVertices.get(v2).add(v1);
+    }
+
+    void removeEdge(String label1, String label2) {
+        Vertex v1 = new Vertex(label1);
+        Vertex v2 = new Vertex(label2);
+        List<Vertex> eV1 = adjVertices.get(v1);
+        List<Vertex> eV2 = adjVertices.get(v2);
+        if (eV1 != null)
+            eV1.remove(v2);
+        if (eV2 != null)
+            eV2.remove(v1);
+    }
+
+    List<Vertex> getAdjVertices(String label) {
+        return adjVertices.get(new Vertex(label));
+    }
+    
+    String printGraph() {
+        StringBuffer sb = new StringBuffer();
+        for(Vertex v : adjVertices.keySet()) {
+            sb.append(v);
+            sb.append(adjVertices.get(v));
+        }
+        return sb.toString();
     }
 
     class Vertex {
@@ -83,30 +73,4 @@ public class Graph {
             return label;
         }
     }
-
-    class Edge {
-        Vertex v1;
-        Vertex v2;
-        Edge(String label1, String label2) {
-            this.v1 = new Vertex(label1);
-            this.v2 = new Vertex(label2);
-        }
-        Edge(Vertex vertex1, Vertex vertex2) {
-            this.v1 = vertex1;
-            this.v2 = vertex2;
-        }
-        @Override
-        public boolean equals(Object obj) {
-            Edge edge = (Edge) obj;
-            return edge.v1.equals(v1) && edge.v2.equals(v2);
-        }
-        @Override
-        public int hashCode() {
-            return v1.hashCode() + v2.hashCode();
-        }
-        @Override
-        public String toString() {
-            return v1.label + "-" + v2.label;
-        }
-    }
 }
\ No newline at end of file
diff --git a/core-java/src/test/java/com/baeldung/graph/GraphTraversalUnitTest.java b/core-java/src/test/java/com/baeldung/graph/GraphTraversalUnitTest.java
index 84611a580a..d955d56d95 100644
--- a/core-java/src/test/java/com/baeldung/graph/GraphTraversalUnitTest.java
+++ b/core-java/src/test/java/com/baeldung/graph/GraphTraversalUnitTest.java
@@ -7,30 +7,30 @@ public class GraphTraversalUnitTest {
     @Test
     public void givenAGraph_whenTraversingDepthFirst_thenExpectedResult() {
         Graph graph = createGraph();
-        Assert.assertEquals("[A, D, E, B, C]", 
-            GraphTraversal.depthFirstTraversal(graph, "A").toString());
+        Assert.assertEquals("[Bob, Rob, Maria, Alice, Mark]", 
+            GraphTraversal.depthFirstTraversal(graph, "Bob").toString());
     }
     
     @Test
     public void givenAGraph_whenTraversingBreadthFirst_thenExpectedResult() {
         Graph graph = createGraph();
-        Assert.assertEquals("[A, B, D, C, E]", 
-            GraphTraversal.breadthFirstTraversal(graph, "A").toString());
+        Assert.assertEquals("[Bob, Alice, Rob, Mark, Maria]", 
+            GraphTraversal.breadthFirstTraversal(graph, "Bob").toString());
     }
     
     Graph createGraph() {
         Graph graph = new Graph();
-        graph.addVertex("A");
-        graph.addVertex("B");
-        graph.addVertex("C");
-        graph.addVertex("D");
-        graph.addVertex("E");
-        graph.addEdge("A", "B");
-        graph.addEdge("A", "D");
-        graph.addEdge("B", "C");
-        graph.addEdge("D", "C");
-        graph.addEdge("B", "E");
-        graph.addEdge("D", "E"); 
+        graph.addVertex("Bob");
+        graph.addVertex("Alice");
+        graph.addVertex("Mark");
+        graph.addVertex("Rob");
+        graph.addVertex("Maria");
+        graph.addEdge("Bob", "Alice");
+        graph.addEdge("Bob", "Rob");
+        graph.addEdge("Alice", "Mark");
+        graph.addEdge("Rob", "Mark");
+        graph.addEdge("Alice", "Maria");
+        graph.addEdge("Rob", "Maria");
         return graph;
     }
 }
\ No newline at end of file