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 12package treemap 13 14import ( 15 "fmt" 16 "github.com/emirpasic/gods/maps" 17 rbt "github.com/emirpasic/gods/trees/redblacktree" 18 "github.com/emirpasic/gods/utils" 19 "strings" 20) 21 22func assertMapImplementation() { 23 var _ maps.Map = (*Map)(nil) 24} 25 26// Map holds the elements in a red-black tree 27type Map struct { 28 tree *rbt.Tree 29} 30 31// NewWith instantiates a tree map with the custom comparator. 32func 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. 37func 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. 42func 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. 48func (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. 55func (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. 61func (m *Map) Remove(key interface{}) { 62 m.tree.Remove(key) 63} 64 65// Empty returns true if map does not contain any elements 66func (m *Map) Empty() bool { 67 return m.tree.Empty() 68} 69 70// Size returns number of elements in the map. 71func (m *Map) Size() int { 72 return m.tree.Size() 73} 74 75// Keys returns all keys in-order 76func (m *Map) Keys() []interface{} { 77 return m.tree.Keys() 78} 79 80// Values returns all values in-order based on the key. 81func (m *Map) Values() []interface{} { 82 return m.tree.Values() 83} 84 85// Clear removes all elements from the map. 86func (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. 92func (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. 101func (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// Floor finds the floor key-value pair for the input key. 109// In case that no floor is found, then both returned values will be nil. 110// It's generally enough to check the first value (key) for nil, which determines if floor was found. 111// 112// Floor key is defined as the largest key that is smaller than or equal to the given key. 113// A floor key may not be found, either because the map is empty, or because 114// all keys in the map are larger than the given key. 115// 116// Key should adhere to the comparator's type assertion, otherwise method panics. 117func (m *Map) Floor(key interface{}) (foundKey interface{}, foundValue interface{}) { 118 node, found := m.tree.Floor(key) 119 if found { 120 return node.Key, node.Value 121 } 122 return nil, nil 123} 124 125// Ceiling finds the ceiling key-value pair for the input key. 126// In case that no ceiling is found, then both returned values will be nil. 127// It's generally enough to check the first value (key) for nil, which determines if ceiling was found. 128// 129// Ceiling key is defined as the smallest key that is larger than or equal to the given key. 130// A ceiling key may not be found, either because the map is empty, or because 131// all keys in the map are smaller than the given key. 132// 133// Key should adhere to the comparator's type assertion, otherwise method panics. 134func (m *Map) Ceiling(key interface{}) (foundKey interface{}, foundValue interface{}) { 135 node, found := m.tree.Ceiling(key) 136 if found { 137 return node.Key, node.Value 138 } 139 return nil, nil 140} 141 142// String returns a string representation of container 143func (m *Map) String() string { 144 str := "TreeMap\nmap[" 145 it := m.Iterator() 146 for it.Next() { 147 str += fmt.Sprintf("%v:%v ", it.Key(), it.Value()) 148 } 149 return strings.TrimRight(str, " ") + "]" 150 151} 152