123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232 |
- // Copyright ©2015 The Gonum Authors. All rights reserved.
- // Use of this source code is governed by a BSD-style
- // license that can be found in the LICENSE file.
- package traverse
- import (
- "gonum.org/v1/gonum/graph"
- "gonum.org/v1/gonum/graph/internal/linear"
- "gonum.org/v1/gonum/graph/internal/set"
- )
- var _ Graph = graph.Graph(nil)
- // Graph is the subset of graph.Graph necessary for graph traversal.
- type Graph interface {
- // From returns all nodes that can be reached directly
- // from the node with the given ID.
- From(id int64) graph.Nodes
- // Edge returns the edge from u to v, with IDs uid and vid,
- // if such an edge exists and nil otherwise. The node v
- // must be directly reachable from u as defined by
- // the From method.
- Edge(uid, vid int64) graph.Edge
- }
- // BreadthFirst implements stateful breadth-first graph traversal.
- type BreadthFirst struct {
- // Visit is called on all nodes on their first visit.
- Visit func(graph.Node)
- // Traverse is called on all edges that may be traversed
- // during the walk. This includes edges that would hop to
- // an already visited node.
- //
- // The value returned by Traverse determines whether
- // an edge can be traversed during the walk.
- Traverse func(graph.Edge) bool
- queue linear.NodeQueue
- visited set.Int64s
- }
- // Walk performs a breadth-first traversal of the graph g starting from the given node,
- // depending on the Traverse field and the until parameter if they are non-nil.
- // The traversal follows edges for which Traverse(edge) is true and returns the first node
- // for which until(node, depth) is true. During the traversal, if the Visit field is
- // non-nil, it is called with each node the first time it is visited.
- func (b *BreadthFirst) Walk(g Graph, from graph.Node, until func(n graph.Node, d int) bool) graph.Node {
- if b.visited == nil {
- b.visited = make(set.Int64s)
- }
- b.queue.Enqueue(from)
- if b.Visit != nil && !b.visited.Has(from.ID()) {
- b.Visit(from)
- }
- b.visited.Add(from.ID())
- var (
- depth int
- children int
- untilNext = 1
- )
- for b.queue.Len() > 0 {
- t := b.queue.Dequeue()
- if until != nil && until(t, depth) {
- return t
- }
- tid := t.ID()
- to := g.From(tid)
- for to.Next() {
- n := to.Node()
- nid := n.ID()
- if b.Traverse != nil && !b.Traverse(g.Edge(tid, nid)) {
- continue
- }
- if b.visited.Has(nid) {
- continue
- }
- if b.Visit != nil {
- b.Visit(n)
- }
- b.visited.Add(nid)
- children++
- b.queue.Enqueue(n)
- }
- if untilNext--; untilNext == 0 {
- depth++
- untilNext = children
- children = 0
- }
- }
- return nil
- }
- // WalkAll calls Walk for each unvisited node of the graph g using edges independent
- // of their direction. The functions before and after are called prior to commencing
- // and after completing each walk if they are non-nil respectively. The function
- // during is called on each node as it is traversed.
- func (b *BreadthFirst) WalkAll(g graph.Undirected, before, after func(), during func(graph.Node)) {
- b.Reset()
- nodes := g.Nodes()
- for nodes.Next() {
- from := nodes.Node()
- if b.Visited(from) {
- continue
- }
- if before != nil {
- before()
- }
- b.Walk(g, from, func(n graph.Node, _ int) bool {
- if during != nil {
- during(n)
- }
- return false
- })
- if after != nil {
- after()
- }
- }
- }
- // Visited returned whether the node n was visited during a traverse.
- func (b *BreadthFirst) Visited(n graph.Node) bool {
- return b.visited.Has(n.ID())
- }
- // Reset resets the state of the traverser for reuse.
- func (b *BreadthFirst) Reset() {
- b.queue.Reset()
- b.visited = nil
- }
- // DepthFirst implements stateful depth-first graph traversal.
- type DepthFirst struct {
- // Visit is called on all nodes on their first visit.
- Visit func(graph.Node)
- // Traverse is called on all edges that may be traversed
- // during the walk. This includes edges that would hop to
- // an already visited node.
- //
- // The value returned by Traverse determines whether an
- // edge can be traversed during the walk.
- Traverse func(graph.Edge) bool
- stack linear.NodeStack
- visited set.Int64s
- }
- // Walk performs a depth-first traversal of the graph g starting from the given node,
- // depending on the Traverse field and the until parameter if they are non-nil.
- // The traversal follows edges for which Traverse(edge) is true and returns the first node
- // for which until(node) is true. During the traversal, if the Visit field is non-nil, it
- // is called with each node the first time it is visited.
- func (d *DepthFirst) Walk(g Graph, from graph.Node, until func(graph.Node) bool) graph.Node {
- if d.visited == nil {
- d.visited = make(set.Int64s)
- }
- d.stack.Push(from)
- if d.Visit != nil && !d.visited.Has(from.ID()) {
- d.Visit(from)
- }
- d.visited.Add(from.ID())
- for d.stack.Len() > 0 {
- t := d.stack.Pop()
- if until != nil && until(t) {
- return t
- }
- tid := t.ID()
- to := g.From(tid)
- for to.Next() {
- n := to.Node()
- nid := n.ID()
- if d.Traverse != nil && !d.Traverse(g.Edge(tid, nid)) {
- continue
- }
- if d.visited.Has(nid) {
- continue
- }
- if d.Visit != nil {
- d.Visit(n)
- }
- d.visited.Add(nid)
- d.stack.Push(n)
- }
- }
- return nil
- }
- // WalkAll calls Walk for each unvisited node of the graph g using edges independent
- // of their direction. The functions before and after are called prior to commencing
- // and after completing each walk if they are non-nil respectively. The function
- // during is called on each node as it is traversed.
- func (d *DepthFirst) WalkAll(g graph.Undirected, before, after func(), during func(graph.Node)) {
- d.Reset()
- nodes := g.Nodes()
- for nodes.Next() {
- from := nodes.Node()
- if d.Visited(from) {
- continue
- }
- if before != nil {
- before()
- }
- d.Walk(g, from, func(n graph.Node) bool {
- if during != nil {
- during(n)
- }
- return false
- })
- if after != nil {
- after()
- }
- }
- }
- // Visited returned whether the node n was visited during a traverse.
- func (d *DepthFirst) Visited(n graph.Node) bool {
- return d.visited.Has(n.ID())
- }
- // Reset resets the state of the traverser for reuse.
- func (d *DepthFirst) Reset() {
- d.stack = d.stack[:0]
- d.visited = nil
- }
|