Graphs

// Package graphs is Module 9: graphs.
//
// Vertices are ints 0..n-1 and edges live in adjacency lists, which is the
// representation you want for almost everything (an adjacency matrix only
// pays off for dense graphs). Neighbors are visited in the order edges were
// added, so traversal orders in the tests are deterministic.
package graphs

// Graph is an unweighted graph stored as adjacency lists.
type Graph struct {
	adj      [][]int
	directed bool
}

// NewGraph returns a graph with n vertices and no edges.
func NewGraph(n int, directed bool) *Graph {
	return &Graph{adj: make([][]int, n), directed: directed}
}

// AddEdge adds u->v (and v->u when the graph is undirected).
func (g *Graph) AddEdge(u, v int) {
	g.adj[u] = append(g.adj[u], v)
	if !g.directed {
		g.adj[v] = append(g.adj[v], u)
	}
}

// N returns the number of vertices.
func (g *Graph) N() int { return len(g.adj) }

// Neighbors returns u's adjacency list (don't modify it).
func (g *Graph) Neighbors(u int) []int { return g.adj[u] }

// Edge is a weighted, directed edge.
type Edge struct {
	To, Weight int
}

// WeightedGraph is a directed graph with non-negative int weights.
type WeightedGraph struct {
	adj [][]Edge
}

// NewWeightedGraph returns a weighted graph with n vertices and no edges.
func NewWeightedGraph(n int) *WeightedGraph {
	return &WeightedGraph{adj: make([][]Edge, n)}
}

// AddEdge adds a directed edge u->v with weight w.
func (g *WeightedGraph) AddEdge(u, v, w int) {
	g.adj[u] = append(g.adj[u], Edge{To: v, Weight: w})
}

// N returns the number of vertices.
func (g *WeightedGraph) N() int { return len(g.adj) }

// Neighbors returns u's outgoing edges (don't modify them).
func (g *WeightedGraph) Neighbors(u int) []Edge { return g.adj[u] }

// UnionFind (disjoint set union) tracks which elements are in the same set.
// parent[i] == i means i is a root. count is the number of distinct sets.
type UnionFind struct {
	parent []int
	rank   []int
	count  int
}

// NewUnionFind returns n singleton sets {0} {1} ... {n-1}.
func NewUnionFind(n int) *UnionFind {
	uf := &UnionFind{parent: make([]int, n), rank: make([]int, n), count: n}
	for i := range uf.parent {
		uf.parent[i] = i
	}
	return uf
}

// Count returns the number of disjoint sets.
func (uf *UnionFind) Count() int { return uf.count }

package graphs

import (
	"slices"
	"testing"
)

// sample is this undirected graph:
//
//	0 — 1 — 3
//	|   |
//	2 — 4    5 — 6
func sample() *Graph {
	g := NewGraph(7, false)
	for _, e := range [][2]int{{0, 1}, {0, 2}, {1, 3}, {1, 4}, {2, 4}, {5, 6}} {
		g.AddEdge(e[0], e[1])
	}
	return g
}

func TestBFS(t *testing.T) {
	if got, want := BFS(sample(), 0), []int{0, 1, 2, 3, 4}; !slices.Equal(got, want) {
		t.Errorf("BFS(0) = %v, want %v", got, want)
	}
	if got, want := BFS(sample(), 5), []int{5, 6}; !slices.Equal(got, want) {
		t.Errorf("BFS(5) = %v, want %v", got, want)
	}
}

func TestDFS(t *testing.T) {
	if got, want := DFS(sample(), 0), []int{0, 1, 3, 4, 2}; !slices.Equal(got, want) {
		t.Errorf("DFS(0) = %v, want %v", got, want)
	}
	if got, want := DFS(sample(), 6), []int{6, 5}; !slices.Equal(got, want) {
		t.Errorf("DFS(6) = %v, want %v", got, want)
	}
}

func TestShortestPath(t *testing.T) {
	g := sample()
	for _, tc := range []struct {
		src, dst int
		wantLen  int // number of vertices on the path; 0 = unreachable
	}{
		{0, 3, 3}, {3, 2, 4}, {2, 2, 1}, {0, 6, 0},
	} {
		path := ShortestPath(g, tc.src, tc.dst)
		if tc.wantLen == 0 {
			if path != nil {
				t.Errorf("ShortestPath(%d, %d) = %v, want nil", tc.src, tc.dst, path)
			}
			continue
		}
		if len(path) != tc.wantLen || path[0] != tc.src || path[len(path)-1] != tc.dst {
			t.Errorf("ShortestPath(%d, %d) = %v, want %d vertices from %d to %d", tc.src, tc.dst, path, tc.wantLen, tc.src, tc.dst)
			continue
		}
		for i := 1; i < len(path); i++ {
			if !slices.Contains(g.Neighbors(path[i-1]), path[i]) {
				t.Errorf("ShortestPath(%d, %d) = %v: %d-%d is not an edge", tc.src, tc.dst, path, path[i-1], path[i])
			}
		}
	}
}

func TestTopoSort(t *testing.T) {
	// Course prerequisites: an edge u->v means "take u before v".
	g := NewGraph(6, true)
	edges := [][2]int{{5, 2}, {5, 0}, {4, 0}, {4, 1}, {2, 3}, {3, 1}}
	for _, e := range edges {
		g.AddEdge(e[0], e[1])
	}
	order, ok := TopoSort(g)
	if !ok || len(order) != 6 {
		t.Fatalf("TopoSort = (%v, %v), want a full order and true", order, ok)
	}
	pos := make(map[int]int)
	for i, v := range order {
		pos[v] = i
	}
	if len(pos) != 6 {
		t.Fatalf("TopoSort = %v: every vertex must appear exactly once", order)
	}
	for _, e := range edges {
		if pos[e[0]] > pos[e[1]] {
			t.Errorf("TopoSort = %v: edge %d->%d is violated", order, e[0], e[1])
		}
	}

	cyclic := NewGraph(3, true)
	cyclic.AddEdge(0, 1)
	cyclic.AddEdge(1, 2)
	cyclic.AddEdge(2, 0)
	if order, ok := TopoSort(cyclic); ok {
		t.Errorf("TopoSort on a cycle = (%v, true), want (nil, false)", order)
	}
}

func TestDijkstra(t *testing.T) {
	g := NewWeightedGraph(6)
	for _, e := range [][3]int{
		{0, 1, 7}, {0, 2, 9}, {0, 5, 14},
		{1, 2, 10}, {1, 3, 15},
		{2, 3, 11}, {2, 5, 2},
		{5, 4, 9}, {3, 4, 6},
	} {
		g.AddEdge(e[0], e[1], e[2])
	}
	if got, want := Dijkstra(g, 0), []int{0, 7, 9, 20, 20, 11}; !slices.Equal(got, want) {
		t.Errorf("Dijkstra(0) = %v, want %v", got, want)
	}
	if got, want := Dijkstra(g, 3), []int{-1, -1, -1, 0, 6, -1}; !slices.Equal(got, want) {
		t.Errorf("Dijkstra(3) = %v, want %v", got, want)
	}
}

func grid(rows ...string) [][]byte {
	out := make([][]byte, len(rows))
	for i, r := range rows {
		out[i] = []byte(r)
	}
	return out
}

func TestNumIslands(t *testing.T) {
	for _, tc := range []struct {
		g    [][]byte
		want int
	}{
		{grid("11110", "11010", "11000", "00000"), 1},
		{grid("11000", "11000", "00100", "00011"), 3},
		{grid("101", "010", "101"), 5},
		{grid("000"), 0},
		{nil, 0},
	} {
		before := slices.Clone(tc.g)
		for i := range before {
			before[i] = slices.Clone(tc.g[i])
		}
		if got := NumIslands(tc.g); got != tc.want {
			t.Errorf("NumIslands(%q) = %d, want %d", before, got, tc.want)
		}
		if !slices.EqualFunc(before, tc.g, slices.Equal[[]byte]) {
			t.Errorf("NumIslands modified its input grid")
		}
	}
}

func TestUnionFind(t *testing.T) {
	uf := NewUnionFind(6)
	if !uf.Union(0, 1) || !uf.Union(2, 3) || !uf.Union(1, 3) {
		t.Fatal("Union of separate sets should return true")
	}
	if uf.Union(0, 2) {
		t.Error("Union(0, 2) should return false: already connected via 1-3")
	}
	if uf.Find(0) != uf.Find(3) {
		t.Error("0 and 3 should share a root")
	}
	if uf.Find(4) == uf.Find(0) {
		t.Error("4 should not be connected to 0")
	}
	if uf.Count() != 3 {
		t.Errorf("Count() = %d, want 3 ({0,1,2,3} {4} {5})", uf.Count())
	}
}

func TestCountComponents(t *testing.T) {
	for _, tc := range []struct {
		n     int
		edges [][2]int
		want  int
	}{
		{5, [][2]int{{0, 1}, {1, 2}, {3, 4}}, 2},
		{5, [][2]int{{0, 1}, {1, 2}, {2, 3}, {3, 4}}, 1},
		{4, nil, 4},
		{3, [][2]int{{0, 1}, {1, 0}, {0, 1}}, 2},
	} {
		if got := CountComponents(tc.n, tc.edges); got != tc.want {
			t.Errorf("CountComponents(%d, %v) = %d, want %d", tc.n, tc.edges, got, tc.want)
		}
	}
}

//go:build solutions

package graphs

import (
	"container/heap"
	"slices"
)

func BFS(g *Graph, start int) []int {
	visited := make([]bool, g.N())
	visited[start] = true
	queue := []int{start}
	var order []int
	for len(queue) > 0 {
		u := queue[0]
		queue = queue[1:]
		order = append(order, u)
		for _, v := range g.Neighbors(u) {
			if !visited[v] {
				visited[v] = true // mark on enqueue, not dequeue
				queue = append(queue, v)
			}
		}
	}
	return order
}

func DFS(g *Graph, start int) []int {
	visited := make([]bool, g.N())
	var order []int
	var visit func(int)
	visit = func(u int) {
		visited[u] = true
		order = append(order, u)
		for _, v := range g.Neighbors(u) {
			if !visited[v] {
				visit(v)
			}
		}
	}
	visit(start)
	return order
}

func ShortestPath(g *Graph, src, dst int) []int {
	parent := make([]int, g.N())
	for i := range parent {
		parent[i] = -1
	}
	visited := make([]bool, g.N())
	visited[src] = true
	queue := []int{src}
	for len(queue) > 0 && !visited[dst] {
		u := queue[0]
		queue = queue[1:]
		for _, v := range g.Neighbors(u) {
			if !visited[v] {
				visited[v] = true
				parent[v] = u
				queue = append(queue, v)
			}
		}
	}
	if !visited[dst] {
		return nil
	}
	var path []int
	for v := dst; v != -1; v = parent[v] {
		path = append(path, v)
	}
	slices.Reverse(path)
	return path
}

func TopoSort(g *Graph) ([]int, bool) {
	indeg := make([]int, g.N())
	for u := 0; u < g.N(); u++ {
		for _, v := range g.Neighbors(u) {
			indeg[v]++
		}
	}
	var queue, order []int
	for u, d := range indeg {
		if d == 0 {
			queue = append(queue, u)
		}
	}
	for len(queue) > 0 {
		u := queue[0]
		queue = queue[1:]
		order = append(order, u)
		for _, v := range g.Neighbors(u) {
			indeg[v]--
			if indeg[v] == 0 {
				queue = append(queue, v)
			}
		}
	}
	if len(order) != g.N() {
		return nil, false // the leftover vertices are on a cycle
	}
	return order, true
}

type item struct{ v, dist int }
type pq []item

func (h pq) Len() int           { return len(h) }
func (h pq) Less(i, j int) bool { return h[i].dist < h[j].dist }
func (h pq) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }
func (h *pq) Push(x any)        { *h = append(*h, x.(item)) }
func (h *pq) Pop() any {
	old := *h
	v := old[len(old)-1]
	*h = old[:len(old)-1]
	return v
}

func Dijkstra(g *WeightedGraph, src int) []int {
	dist := make([]int, g.N())
	for i := range dist {
		dist[i] = -1
	}
	dist[src] = 0
	h := &pq{{src, 0}}
	for h.Len() > 0 {
		cur := heap.Pop(h).(item)
		if cur.dist > dist[cur.v] {
			continue // stale entry: we already found a shorter path
		}
		for _, e := range g.Neighbors(cur.v) {
			nd := cur.dist + e.Weight
			if dist[e.To] == -1 || nd < dist[e.To] {
				dist[e.To] = nd
				heap.Push(h, item{e.To, nd})
			}
		}
	}
	return dist
}

func NumIslands(grid [][]byte) int {
	rows := len(grid)
	if rows == 0 {
		return 0
	}
	cols := len(grid[0])
	seen := make([][]bool, rows)
	for i := range seen {
		seen[i] = make([]bool, cols)
	}
	var sink func(r, c int)
	sink = func(r, c int) {
		if r < 0 || r >= rows || c < 0 || c >= cols || seen[r][c] || grid[r][c] != '1' {
			return
		}
		seen[r][c] = true
		sink(r+1, c)
		sink(r-1, c)
		sink(r, c+1)
		sink(r, c-1)
	}
	count := 0
	for r := 0; r < rows; r++ {
		for c := 0; c < cols; c++ {
			if grid[r][c] == '1' && !seen[r][c] {
				count++
				sink(r, c)
			}
		}
	}
	return count
}

func (uf *UnionFind) Find(x int) int {
	for uf.parent[x] != x {
		uf.parent[x] = uf.parent[uf.parent[x]] // path halving
		x = uf.parent[x]
	}
	return x
}

func (uf *UnionFind) Union(a, b int) bool {
	ra, rb := uf.Find(a), uf.Find(b)
	if ra == rb {
		return false
	}
	switch {
	case uf.rank[ra] < uf.rank[rb]:
		uf.parent[ra] = rb
	case uf.rank[ra] > uf.rank[rb]:
		uf.parent[rb] = ra
	default:
		uf.parent[rb] = ra
		uf.rank[ra]++
	}
	uf.count--
	return true
}

func CountComponents(n int, edges [][2]int) int {
	uf := NewUnionFind(n)
	for _, e := range edges {
		uf.Union(e[0], e[1])
	}
	return uf.Count()
}