go: add «841. Keys and Rooms»

Signed-off-by: Matej Focko <me@mfocko.xyz>

go/keys-and-rooms.go

@@ -0,0 +1,48 @@
+package main
+
+import (
+	arraystack "github.com/emirpasic/gods/v2/stacks/arraystack"
+)
+
+type DFSEntry struct {
+	room  int
+	index int
+}
+
+func (e *DFSEntry) Next() DFSEntry {
+	return DFSEntry{room: e.room, index: e.index + 1}
+}
+
+func canVisitAllRooms(rooms [][]int) bool {
+	s := arraystack.New[DFSEntry]()
+	s.Push(DFSEntry{room: 0, index: 0})
+
+	visited := make([]bool, len(rooms))
+	visited[0] = true
+
+	for !s.Empty() {
+		key, _ := s.Pop()
+
+		if key.index >= len(rooms[key.room]) {
+			continue
+		}
+
+		s.Push(key.Next())
+
+		unlockedRoom := rooms[key.room][key.index]
+		if visited[unlockedRoom] {
+			continue
+		}
+
+		visited[unlockedRoom] = true
+		s.Push(DFSEntry{room: unlockedRoom, index: 0})
+	}
+
+	for _, unlocked := range visited {
+		if !unlocked {
+			return false
+		}
+	}
+
+	return true
+}