traverse.go 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232
  1. // Copyright ©2015 The Gonum Authors. All rights reserved.
  2. // Use of this source code is governed by a BSD-style
  3. // license that can be found in the LICENSE file.
  4. package traverse
  5. import (
  6. "gonum.org/v1/gonum/graph"
  7. "gonum.org/v1/gonum/graph/internal/linear"
  8. "gonum.org/v1/gonum/graph/internal/set"
  9. )
  10. var _ Graph = graph.Graph(nil)
  11. // Graph is the subset of graph.Graph necessary for graph traversal.
  12. type Graph interface {
  13. // From returns all nodes that can be reached directly
  14. // from the node with the given ID.
  15. From(id int64) graph.Nodes
  16. // Edge returns the edge from u to v, with IDs uid and vid,
  17. // if such an edge exists and nil otherwise. The node v
  18. // must be directly reachable from u as defined by
  19. // the From method.
  20. Edge(uid, vid int64) graph.Edge
  21. }
  22. // BreadthFirst implements stateful breadth-first graph traversal.
  23. type BreadthFirst struct {
  24. // Visit is called on all nodes on their first visit.
  25. Visit func(graph.Node)
  26. // Traverse is called on all edges that may be traversed
  27. // during the walk. This includes edges that would hop to
  28. // an already visited node.
  29. //
  30. // The value returned by Traverse determines whether
  31. // an edge can be traversed during the walk.
  32. Traverse func(graph.Edge) bool
  33. queue linear.NodeQueue
  34. visited set.Int64s
  35. }
  36. // Walk performs a breadth-first traversal of the graph g starting from the given node,
  37. // depending on the Traverse field and the until parameter if they are non-nil.
  38. // The traversal follows edges for which Traverse(edge) is true and returns the first node
  39. // for which until(node, depth) is true. During the traversal, if the Visit field is
  40. // non-nil, it is called with each node the first time it is visited.
  41. func (b *BreadthFirst) Walk(g Graph, from graph.Node, until func(n graph.Node, d int) bool) graph.Node {
  42. if b.visited == nil {
  43. b.visited = make(set.Int64s)
  44. }
  45. b.queue.Enqueue(from)
  46. if b.Visit != nil && !b.visited.Has(from.ID()) {
  47. b.Visit(from)
  48. }
  49. b.visited.Add(from.ID())
  50. var (
  51. depth int
  52. children int
  53. untilNext = 1
  54. )
  55. for b.queue.Len() > 0 {
  56. t := b.queue.Dequeue()
  57. if until != nil && until(t, depth) {
  58. return t
  59. }
  60. tid := t.ID()
  61. to := g.From(tid)
  62. for to.Next() {
  63. n := to.Node()
  64. nid := n.ID()
  65. if b.Traverse != nil && !b.Traverse(g.Edge(tid, nid)) {
  66. continue
  67. }
  68. if b.visited.Has(nid) {
  69. continue
  70. }
  71. if b.Visit != nil {
  72. b.Visit(n)
  73. }
  74. b.visited.Add(nid)
  75. children++
  76. b.queue.Enqueue(n)
  77. }
  78. if untilNext--; untilNext == 0 {
  79. depth++
  80. untilNext = children
  81. children = 0
  82. }
  83. }
  84. return nil
  85. }
  86. // WalkAll calls Walk for each unvisited node of the graph g using edges independent
  87. // of their direction. The functions before and after are called prior to commencing
  88. // and after completing each walk if they are non-nil respectively. The function
  89. // during is called on each node as it is traversed.
  90. func (b *BreadthFirst) WalkAll(g graph.Undirected, before, after func(), during func(graph.Node)) {
  91. b.Reset()
  92. nodes := g.Nodes()
  93. for nodes.Next() {
  94. from := nodes.Node()
  95. if b.Visited(from) {
  96. continue
  97. }
  98. if before != nil {
  99. before()
  100. }
  101. b.Walk(g, from, func(n graph.Node, _ int) bool {
  102. if during != nil {
  103. during(n)
  104. }
  105. return false
  106. })
  107. if after != nil {
  108. after()
  109. }
  110. }
  111. }
  112. // Visited returned whether the node n was visited during a traverse.
  113. func (b *BreadthFirst) Visited(n graph.Node) bool {
  114. return b.visited.Has(n.ID())
  115. }
  116. // Reset resets the state of the traverser for reuse.
  117. func (b *BreadthFirst) Reset() {
  118. b.queue.Reset()
  119. b.visited = nil
  120. }
  121. // DepthFirst implements stateful depth-first graph traversal.
  122. type DepthFirst struct {
  123. // Visit is called on all nodes on their first visit.
  124. Visit func(graph.Node)
  125. // Traverse is called on all edges that may be traversed
  126. // during the walk. This includes edges that would hop to
  127. // an already visited node.
  128. //
  129. // The value returned by Traverse determines whether an
  130. // edge can be traversed during the walk.
  131. Traverse func(graph.Edge) bool
  132. stack linear.NodeStack
  133. visited set.Int64s
  134. }
  135. // Walk performs a depth-first traversal of the graph g starting from the given node,
  136. // depending on the Traverse field and the until parameter if they are non-nil.
  137. // The traversal follows edges for which Traverse(edge) is true and returns the first node
  138. // for which until(node) is true. During the traversal, if the Visit field is non-nil, it
  139. // is called with each node the first time it is visited.
  140. func (d *DepthFirst) Walk(g Graph, from graph.Node, until func(graph.Node) bool) graph.Node {
  141. if d.visited == nil {
  142. d.visited = make(set.Int64s)
  143. }
  144. d.stack.Push(from)
  145. if d.Visit != nil && !d.visited.Has(from.ID()) {
  146. d.Visit(from)
  147. }
  148. d.visited.Add(from.ID())
  149. for d.stack.Len() > 0 {
  150. t := d.stack.Pop()
  151. if until != nil && until(t) {
  152. return t
  153. }
  154. tid := t.ID()
  155. to := g.From(tid)
  156. for to.Next() {
  157. n := to.Node()
  158. nid := n.ID()
  159. if d.Traverse != nil && !d.Traverse(g.Edge(tid, nid)) {
  160. continue
  161. }
  162. if d.visited.Has(nid) {
  163. continue
  164. }
  165. if d.Visit != nil {
  166. d.Visit(n)
  167. }
  168. d.visited.Add(nid)
  169. d.stack.Push(n)
  170. }
  171. }
  172. return nil
  173. }
  174. // WalkAll calls Walk for each unvisited node of the graph g using edges independent
  175. // of their direction. The functions before and after are called prior to commencing
  176. // and after completing each walk if they are non-nil respectively. The function
  177. // during is called on each node as it is traversed.
  178. func (d *DepthFirst) WalkAll(g graph.Undirected, before, after func(), during func(graph.Node)) {
  179. d.Reset()
  180. nodes := g.Nodes()
  181. for nodes.Next() {
  182. from := nodes.Node()
  183. if d.Visited(from) {
  184. continue
  185. }
  186. if before != nil {
  187. before()
  188. }
  189. d.Walk(g, from, func(n graph.Node) bool {
  190. if during != nil {
  191. during(n)
  192. }
  193. return false
  194. })
  195. if after != nil {
  196. after()
  197. }
  198. }
  199. }
  200. // Visited returned whether the node n was visited during a traverse.
  201. func (d *DepthFirst) Visited(n graph.Node) bool {
  202. return d.visited.Has(n.ID())
  203. }
  204. // Reset resets the state of the traverser for reuse.
  205. func (d *DepthFirst) Reset() {
  206. d.stack = d.stack[:0]
  207. d.visited = nil
  208. }