Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?

Data structures in Go with maps and slices

Sets:

set := make(map[string]bool)

// add to set
set["a"] = true 

// delete from set
delete(set, "a") 

// set membership
if set["b"] {
  // do stuff
}

// iterate over elements
for t := range set {
  // do stuff
}

Queues:

var q []string

// enqueue
q = append(q, "a")

// dequeue until empty
for len(q) > 0 {
  first, q = q[0], q[1:] 
  // do stuff
}

Stacks:

var stack []string

// push onto stack
stack = append(stack, "a")

// pop until empty
for len(stack) > 0 {
  last, stack = stack[len(stack)-1], stack[:len(stack)-1]
  // do stuff
}

Graphs:

// node -> neighbor set
g := make(map[string]map[string]bool)

// add edges
g["a"] = map[string]int{}
g["a"]["b"] = true
g["a"]["c"] = true
g["b"]["a"] = true

// remove edge
delete(g["a"], "c")

// depth-first search
func dfs(
  g map[string]map[string]bool,
  node string,
  visited map[string]bool,
  visit func(string),
) {
  visit(node)
  visited[node] = true
  for nbr := range g[node] {
    if !visited[nbr] {
      dfs(g, nbr, visited, visit)
    }
  }
}

type edge struct {
  node string
  dist int
}

// breadth-first search
func bfs(
  g map[string]map[string]bool,
  from node,
  visit func(node string, dist int),
) {
  q := []edge{{from, 0}}
  v := make(map[string]bool)
  var first edge
  for len(q) > 0 {
    first, q = q[0], q[1:]
    v[first.node] = true
    visit(first.node, first.dist)
    for nbr := range g[first.from] {
      if !v[nbr] {
        q = append(q, edge{nbr, first.dist+1})
      }
    }
  }
}

type pathEdge struct {
  node string
  path []string
}

// shortest paths with bfs
func shortestPaths(
  g map[string]map[string]bool,
  from node,
) map[string][]string {
  paths := make(map[string][]string)
  q := []pathEdge{{from, []string{from}}}
  v := make(map[string]bool)
  var first pathEdge
  for len(q) > 0 {
    first, q = q[0], q[1:]
    v[first.node] = true
    paths[first.node] = first.path
    for nbr := range g[first.node] {
      if !v[nbr] {
        path := append([]string{}, first.path...)
        q = append(q, pathEdge{nbr, append(path, nbr)})
      }
    }
  }
  return paths
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.