Map in order range loop

by

If you need a map and keys in order, those are 2 different things, you need 2 different (data) types to provide that functionality.

With a keys slice

The easiest way to achieve this is to maintain key order in a different slice. Whenever you put a new pair into the map, first check if the key is already in it. If not, add the new key to the separate slice. When you need elements in order, you may use the keys slice. Of course when you remove a pair, you also have to remove it from the slice too.

The keys slice only has to contain the keys (and not the values), so the overhead is little.

Wrap this new functionality (map+keys slice) into a new type and provide methods for it, and hide the map and slice. Then data misalignment cannot occur.

Example implementation:

type Key int   // Key type
type Value int // Value type

type Map struct {
    m    map[Key]Value
    keys []Key
}

func New() *Map {
    return &Map{m: make(map[Key]Value)}
}

func (m *Map) Set(k Key, v Value) {
    if _, ok := m.m[k]; !ok {
        m.keys = append(m.keys, k)
    }
    m.m[k] = v
}

func (m *Map) Range() {
    for _, k := range m.keys {
        fmt.Println(m.m[k])
    }
}

Using it:

m := New()
m.Set(1, 11)
m.Set(2, 22)
m.Range()

Try it on the Go Playground.

With a value-wrapper implementing a linked-list

Another approach would be to wrap the values, and –along the real value– also store the next/previous key.

For example, assuming you want a map like map[Key]Value:

type valueWrapper struct {
    value Value
    next  *Key // Next key
}

Whenever you add a pair to the map, you set a valueWrapper as the value, and you have to link this to the previous (last) pair. To link, you have to set next field of the last wrapper to point to this new key. To easily implement this, it’s recommended to also store the last key (to avoid having to search for it).

When you want to iterate over the elements in insertion order, you start from the first (you have to store this), and its associated valueWrapper will tell you the next key (in insertion order).

Example implementation:

type Key int   // Key type
type Value int // Value type

type valueWrapper struct {
    v    Value
    next *Key
}

type Map struct {
    m           map[Key]valueWrapper
    first, last *Key
}

func New() *Map {
    return &Map{m: make(map[Key]valueWrapper)}
}

func (m *Map) Set(k Key, v Value) {
    if _, ok := m.m[k]; !ok && m.last != nil {
        w2 := m.m[*m.last]
        m.m[*m.last] = valueWrapper{w2.v, &k}
    }
    w := valueWrapper{v: v}
    m.m[k] = w
    if m.first == nil {
        m.first = &k
    }
    m.last = &k
}

func (m *Map) Range() {
    for k := m.first; k != nil; {
        w := m.m[*k]
        fmt.Println(w.v)
        k = w.next
    }
}

Using it is the same. Try it on the Go Playground.

Notes: You may vary a couple of things to your liking:

  • You may declare the internal map like m map[Key]*valueWrapper and so in Set() you can change the next field without having to assign a new valueWrapper.

  • You may choose first and last fields to be of type *valueWrapper

  • You may choose next to be of type *valueWrapper

Comparison

The approach with an additional slice is easier and cleaner. But removing an element from it may become slow if the map grows big, as we also have to find the key in the slice which is “unsorted”, so it’s O(n) complexity.

The approach with linked-list in value-wrapper can easily be extended to support fast element removal even if the map is big, if you also add the prev field to the valueWrapper struct. So if you need to remove an element, you can super-fast find the wrapper (O(1)), update the prev and next wrappers (to point to each other), and perform a simple delete() operation, it’s O(1).

Note that deletion in the first solution (with slice) could still be sped up by using 1 additional map, which would map from key to index of the key in the slice (map[Key]int), so delete operation could still be implemented in O(1), in exchange for greater complexity. Another option for speeding up could be to change the value in the map to be a wrapper, which could hold the actual value and the index of the key in the slice.

See related question: Why can’t Go iterate maps in insertion order?

Leave a Comment