github.com/alexandrestein/gods@v1.0.1/maps/treemap/treemap.go (about)

     1  // Copyright (c) 2015, Emir Pasic. 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  
     5  // Package treemap implements a map backed by red-black tree.
     6  //
     7  // Elements are ordered by key in the map.
     8  //
     9  // Structure is not thread safe.
    10  //
    11  // Reference: http://en.wikipedia.org/wiki/Associative_array
    12  package treemap
    13  
    14  import (
    15  	"fmt"
    16  	"github.com/alexandrestein/gods/maps"
    17  	rbt "github.com/alexandrestein/gods/trees/redblacktree"
    18  	"github.com/alexandrestein/gods/utils"
    19  	"strings"
    20  )
    21  
    22  func assertMapImplementation() {
    23  	var _ maps.Map = (*Map)(nil)
    24  }
    25  
    26  // Map holds the elements in a red-black tree
    27  type Map struct {
    28  	tree *rbt.Tree
    29  }
    30  
    31  // NewWith instantiates a tree map with the custom comparator.
    32  func NewWith(comparator utils.Comparator) *Map {
    33  	return &Map{tree: rbt.NewWith(comparator)}
    34  }
    35  
    36  // NewWithIntComparator instantiates a tree map with the IntComparator, i.e. keys are of type int.
    37  func NewWithIntComparator() *Map {
    38  	return &Map{tree: rbt.NewWithIntComparator()}
    39  }
    40  
    41  // NewWithStringComparator instantiates a tree map with the StringComparator, i.e. keys are of type string.
    42  func NewWithStringComparator() *Map {
    43  	return &Map{tree: rbt.NewWithStringComparator()}
    44  }
    45  
    46  // Put inserts key-value pair into the map.
    47  // Key should adhere to the comparator's type assertion, otherwise method panics.
    48  func (m *Map) Put(key interface{}, value interface{}) {
    49  	m.tree.Put(key, value)
    50  }
    51  
    52  // Get searches the element in the map by key and returns its value or nil if key is not found in tree.
    53  // Second return parameter is true if key was found, otherwise false.
    54  // Key should adhere to the comparator's type assertion, otherwise method panics.
    55  func (m *Map) Get(key interface{}) (value interface{}, found bool) {
    56  	return m.tree.Get(key)
    57  }
    58  
    59  // Remove removes the element from the map by key.
    60  // Key should adhere to the comparator's type assertion, otherwise method panics.
    61  func (m *Map) Remove(key interface{}) {
    62  	m.tree.Remove(key)
    63  }
    64  
    65  // Empty returns true if map does not contain any elements
    66  func (m *Map) Empty() bool {
    67  	return m.tree.Empty()
    68  }
    69  
    70  // Size returns number of elements in the map.
    71  func (m *Map) Size() int {
    72  	return m.tree.Size()
    73  }
    74  
    75  // Keys returns all keys in-order
    76  func (m *Map) Keys() []interface{} {
    77  	return m.tree.Keys()
    78  }
    79  
    80  // Values returns all values in-order based on the key.
    81  func (m *Map) Values() []interface{} {
    82  	return m.tree.Values()
    83  }
    84  
    85  // Clear removes all elements from the map.
    86  func (m *Map) Clear() {
    87  	m.tree.Clear()
    88  }
    89  
    90  // Min returns the minimum key and its value from the tree map.
    91  // Returns nil, nil if map is empty.
    92  func (m *Map) Min() (key interface{}, value interface{}) {
    93  	if node := m.tree.Left(); node != nil {
    94  		return node.Key, node.Value
    95  	}
    96  	return nil, nil
    97  }
    98  
    99  // Max returns the maximum key and its value from the tree map.
   100  // Returns nil, nil if map is empty.
   101  func (m *Map) Max() (key interface{}, value interface{}) {
   102  	if node := m.tree.Right(); node != nil {
   103  		return node.Key, node.Value
   104  	}
   105  	return nil, nil
   106  }
   107  
   108  // String returns a string representation of container
   109  func (m *Map) String() string {
   110  	str := "TreeMap\nmap["
   111  	it := m.Iterator()
   112  	for it.Next() {
   113  		str += fmt.Sprintf("%v:%v ", it.Key(), it.Value())
   114  	}
   115  	return strings.TrimRight(str, " ") + "]"
   116  
   117  }