graph.go

// Package graph implements directed and weighted directed graphs with
// adjacency-list storage, plus BFS, DFS and topological sort.
package graph

// Graph is a directed graph with integer vertices in [0, n) and unweighted
// edges, stored as an adjacency list.
type Graph struct {
	adj [][]int
}

// New returns a new Graph with n vertices and no edges.
func New(n int) *Graph {
	if n < 0 {
		n = 0
	}
	return &Graph{adj: make([][]int, n)}
}

// FromEdges constructs a Graph with n vertices and adds the given edges.
func FromEdges(n int, edges [][2]int) *Graph {
	g := New(n)
	for _, e := range edges {
		g.AddEdge(e[0], e[1])
	}
	return g
}

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

// AddEdge adds a directed edge from u to v.
func (g *Graph) AddEdge(u, v int) {
	g.adj[u] = append(g.adj[u], v)
}

// Neighbors returns the (un-copied) slice of out-neighbors of u. Do not
// modify the returned slice.
func (g *Graph) Neighbors(u int) []int { return g.adj[u] }

// BFS visits each vertex reachable from src in breadth-first order.
func (g *Graph) BFS(src int, visit func(int)) {
	n := len(g.adj)
	if src < 0 || src >= n {
		return
	}
	seen := make([]bool, n)
	q := make([]int, 0, n)
	seen[src] = true
	q = append(q, src)
	for len(q) > 0 {
		u := q[0]
		q = q[1:]
		visit(u)
		for _, v := range g.adj[u] {
			if !seen[v] {
				seen[v] = true
				q = append(q, v)
			}
		}
	}
}

// DFSIter performs an iterative depth-first search from src. The traversal
// order matches the recursive DFS that processes neighbors in declared order
// (we push them in reverse so the first neighbor is visited first).
func (g *Graph) DFSIter(src int, visit func(int)) {
	n := len(g.adj)
	if src < 0 || src >= n {
		return
	}
	seen := make([]bool, n)
	st := []int{src}
	for len(st) > 0 {
		u := st[len(st)-1]
		st = st[:len(st)-1]
		if seen[u] {
			continue
		}
		seen[u] = true
		visit(u)
		neighbors := g.adj[u]
		for i := len(neighbors) - 1; i >= 0; i-- {
			if !seen[neighbors[i]] {
				st = append(st, neighbors[i])
			}
		}
	}
}

// DFSRec performs a recursive depth-first search from src.
func (g *Graph) DFSRec(src int, visit func(int)) {
	n := len(g.adj)
	if src < 0 || src >= n {
		return
	}
	seen := make([]bool, n)
	var dfs func(int)
	dfs = func(u int) {
		seen[u] = true
		visit(u)
		for _, v := range g.adj[u] {
			if !seen[v] {
				dfs(v)
			}
		}
	}
	dfs(src)
}

// TopologicalSort returns a topological ordering using Kahn's algorithm.
// The bool is false (and the slice nil) if the graph contains a cycle.
//
// Time:  O(V + E).
// Space: O(V).
func (g *Graph) TopologicalSort() ([]int, bool) {
	n := len(g.adj)
	indeg := make([]int, n)
	for u := 0; u < n; u++ {
		for _, v := range g.adj[u] {
			indeg[v]++
		}
	}
	q := make([]int, 0, n)
	for i := 0; i < n; i++ {
		if indeg[i] == 0 {
			q = append(q, i)
		}
	}
	order := make([]int, 0, n)
	for len(q) > 0 {
		u := q[0]
		q = q[1:]
		order = append(order, u)
		for _, v := range g.adj[u] {
			indeg[v]--
			if indeg[v] == 0 {
				q = append(q, v)
			}
		}
	}
	if len(order) != n {
		return nil, false
	}
	return order, true
}