Kenenbek Arzymatov
Kenenbek Arzymatov

Reputation: 9109

Construct user defined map in Go

I want to create map which stores routes from one point to another, distance between them will be used as values in the map.

type Route struct {
    start  string
    finish string
}

m := make(map[Route]int)
v := Route{start: "A",
            finish: "B"}
m[v] = 42

Distance from A to B equals to distance from B to A.

w := Route{start: "B",
        finish: "A"}

How can I get m[w] without pushing m[w] = 42 into the map again because v and w are the same route.

P.S In other words, is it possible to override rules when one Route equals to another Route when using Route as a key in map?

Upvotes: 1

Views: 103

Answers (3)

Ilayaraja
Ilayaraja

Reputation: 2876

Define map[Route]int as a user-defined type like RouteDistance.

type RouteDistance map[Route]int

And then define a Find function on that type, to check for bi-directional routes and return the distance.

func (rd RouteDistance) Find(r Route) (d int, ok bool) {
    routeDistance := map[Route]int(rd)
    d, ok = routeDistance[r]
    if ok {
        return
    }
    d, ok = routeDistance[Route{start: r.finish, finish: r.start}]
    if ok {
        return
    }
    return -1, false
}

Initiate the RouteDistance with just one-directional route and distance.

rd := make(RouteDistance)
rd[Route{start: "A", finish: "B"}] = 42

Get bi-directional distance for routes.

rd.Find(Route{start: "A", finish: "B"})
rd.Find(Route{start: "B", finish: "A"})

The working code is in Playground

Upvotes: 2

abhink
abhink

Reputation: 9116

A few ways to do this without pushing an additional value. Either perform two map accesses:

r := Route{
    start: "B",
    finish: "A",
}
d, ok := m[r]
if !ok {
    r.start = "B"
    r.finish = "A"
    d, ok = m[r]
}
if !ok {
    // value not found
}

Or for every map store and access, sort the start and finish:

func Put(start, finish string, d int) {
    r := Route{start, finish}
    if start > finish {
        r.start = finish
        r.finish = start
    }
    m[r] = d
}

func Get(r Route) int {
    if r.start > r.finish {
        r.start, r.finish = r.finish, r.start
    }
    m[r] = d
}

Upvotes: 1

icza
icza

Reputation: 417512

Index the map with a Route where finish and start are swapped:

m[Route{start: w.finish, finish: w.start}]

If you want to check both directions, create a helper function for it:

var m = make(map[Route]int)

func dist(r Route) int {
    if d, ok := m[r]; ok {
        return d
    }
    return m[Route{start: r.finish, finish: r.start}]
}

Note that dist() will return 0 if neither direction is in the map. If you want to return this info too:

func dist(r Route) (d int, found bool) {
    d, found = m[r]
    if !found {
        d, found = m[Route{start: r.finish, finish: r.start}]
    }
    return
}

Try it on the Go Playground.

Upvotes: 1

Related Questions