1// Copyright ©2014 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
5package simple
6
7import (
8	"fmt"
9
10	"gonum.org/v1/gonum/graph"
11	"gonum.org/v1/gonum/graph/internal/uid"
12	"gonum.org/v1/gonum/graph/iterator"
13)
14
15var (
16	ug *UndirectedGraph
17
18	_ graph.Graph       = ug
19	_ graph.Undirected  = ug
20	_ graph.NodeAdder   = ug
21	_ graph.NodeRemover = ug
22	_ graph.EdgeAdder   = ug
23	_ graph.EdgeRemover = ug
24)
25
26// UndirectedGraph implements a generalized undirected graph.
27type UndirectedGraph struct {
28	nodes map[int64]graph.Node
29	edges map[int64]map[int64]graph.Edge
30
31	nodeIDs uid.Set
32}
33
34// NewUndirectedGraph returns an UndirectedGraph.
35func NewUndirectedGraph() *UndirectedGraph {
36	return &UndirectedGraph{
37		nodes: make(map[int64]graph.Node),
38		edges: make(map[int64]map[int64]graph.Edge),
39
40		nodeIDs: uid.NewSet(),
41	}
42}
43
44// AddNode adds n to the graph. It panics if the added node ID matches an existing node ID.
45func (g *UndirectedGraph) AddNode(n graph.Node) {
46	if _, exists := g.nodes[n.ID()]; exists {
47		panic(fmt.Sprintf("simple: node ID collision: %d", n.ID()))
48	}
49	g.nodes[n.ID()] = n
50	g.edges[n.ID()] = make(map[int64]graph.Edge)
51	g.nodeIDs.Use(n.ID())
52}
53
54// Edge returns the edge from u to v if such an edge exists and nil otherwise.
55// The node v must be directly reachable from u as defined by the From method.
56func (g *UndirectedGraph) Edge(uid, vid int64) graph.Edge {
57	return g.EdgeBetween(uid, vid)
58}
59
60// EdgeBetween returns the edge between nodes x and y.
61func (g *UndirectedGraph) EdgeBetween(xid, yid int64) graph.Edge {
62	edge, ok := g.edges[xid][yid]
63	if !ok {
64		return nil
65	}
66	if edge.From().ID() == xid {
67		return edge
68	}
69	return edge.ReversedEdge()
70}
71
72// Edges returns all the edges in the graph.
73func (g *UndirectedGraph) Edges() graph.Edges {
74	if len(g.edges) == 0 {
75		return graph.Empty
76	}
77	var edges []graph.Edge
78	seen := make(map[[2]int64]struct{})
79	for _, u := range g.edges {
80		for _, e := range u {
81			uid := e.From().ID()
82			vid := e.To().ID()
83			if _, ok := seen[[2]int64{uid, vid}]; ok {
84				continue
85			}
86			seen[[2]int64{uid, vid}] = struct{}{}
87			seen[[2]int64{vid, uid}] = struct{}{}
88			edges = append(edges, e)
89		}
90	}
91	if len(edges) == 0 {
92		return graph.Empty
93	}
94	return iterator.NewOrderedEdges(edges)
95}
96
97// From returns all nodes in g that can be reached directly from n.
98func (g *UndirectedGraph) From(id int64) graph.Nodes {
99	if _, ok := g.nodes[id]; !ok {
100		return graph.Empty
101	}
102
103	nodes := make([]graph.Node, len(g.edges[id]))
104	i := 0
105	for from := range g.edges[id] {
106		nodes[i] = g.nodes[from]
107		i++
108	}
109	if len(nodes) == 0 {
110		return graph.Empty
111	}
112	return iterator.NewOrderedNodes(nodes)
113}
114
115// HasEdgeBetween returns whether an edge exists between nodes x and y.
116func (g *UndirectedGraph) HasEdgeBetween(xid, yid int64) bool {
117	_, ok := g.edges[xid][yid]
118	return ok
119}
120
121// NewEdge returns a new Edge from the source to the destination node.
122func (g *UndirectedGraph) NewEdge(from, to graph.Node) graph.Edge {
123	return &Edge{F: from, T: to}
124}
125
126// NewNode returns a new unique Node to be added to g. The Node's ID does
127// not become valid in g until the Node is added to g.
128func (g *UndirectedGraph) NewNode() graph.Node {
129	if len(g.nodes) == 0 {
130		return Node(0)
131	}
132	if int64(len(g.nodes)) == uid.Max {
133		panic("simple: cannot allocate node: no slot")
134	}
135	return Node(g.nodeIDs.NewID())
136}
137
138// Node returns the node with the given ID if it exists in the graph,
139// and nil otherwise.
140func (g *UndirectedGraph) Node(id int64) graph.Node {
141	return g.nodes[id]
142}
143
144// Nodes returns all the nodes in the graph.
145func (g *UndirectedGraph) Nodes() graph.Nodes {
146	if len(g.nodes) == 0 {
147		return graph.Empty
148	}
149	nodes := make([]graph.Node, len(g.nodes))
150	i := 0
151	for _, n := range g.nodes {
152		nodes[i] = n
153		i++
154	}
155	return iterator.NewOrderedNodes(nodes)
156}
157
158// RemoveEdge removes the edge with the given end IDs from the graph, leaving the terminal nodes.
159// If the edge does not exist it is a no-op.
160func (g *UndirectedGraph) RemoveEdge(fid, tid int64) {
161	if _, ok := g.nodes[fid]; !ok {
162		return
163	}
164	if _, ok := g.nodes[tid]; !ok {
165		return
166	}
167
168	delete(g.edges[fid], tid)
169	delete(g.edges[tid], fid)
170}
171
172// RemoveNode removes the node with the given ID from the graph, as well as any edges attached
173// to it. If the node is not in the graph it is a no-op.
174func (g *UndirectedGraph) RemoveNode(id int64) {
175	if _, ok := g.nodes[id]; !ok {
176		return
177	}
178	delete(g.nodes, id)
179
180	for from := range g.edges[id] {
181		delete(g.edges[from], id)
182	}
183	delete(g.edges, id)
184
185	g.nodeIDs.Release(id)
186}
187
188// SetEdge adds e, an edge from one node to another. If the nodes do not exist, they are added
189// and are set to the nodes of the edge otherwise.
190// It will panic if the IDs of the e.From and e.To are equal.
191func (g *UndirectedGraph) SetEdge(e graph.Edge) {
192	var (
193		from = e.From()
194		fid  = from.ID()
195		to   = e.To()
196		tid  = to.ID()
197	)
198
199	if fid == tid {
200		panic("simple: adding self edge")
201	}
202
203	if _, ok := g.nodes[fid]; !ok {
204		g.AddNode(from)
205	} else {
206		g.nodes[fid] = from
207	}
208	if _, ok := g.nodes[tid]; !ok {
209		g.AddNode(to)
210	} else {
211		g.nodes[tid] = to
212	}
213
214	g.edges[fid][tid] = e
215	g.edges[tid][fid] = e
216}
217