gitee.com/sy_183/go-common@v1.0.5-0.20231205030221-958cfe129b47/container/ordered-tree-map.go (about)

     1  package container
     2  
     3  import "gitee.com/sy_183/go-common/generic"
     4  
     5  type OrderedTreeMap[K generic.Ordered, V any] struct {
     6  	entries []*TreeMapEntry[K, V]
     7  	root    *TreeMapEntry[K, V]
     8  	head    *TreeMapEntry[K, V]
     9  	tail    *TreeMapEntry[K, V]
    10  }
    11  
    12  func NewOrderedTreeMap[K generic.Ordered, V any](entries int) *OrderedTreeMap[K, V] {
    13  	m := &OrderedTreeMap[K, V]{
    14  		entries: make([]*TreeMapEntry[K, V], entries),
    15  	}
    16  	m.init()
    17  	return m
    18  }
    19  
    20  func (m *OrderedTreeMap[K, V]) init() {
    21  	rawEntries := make([]TreeMapEntry[K, V], len(m.entries))
    22  	for i := range rawEntries {
    23  		m.entries[i] = &rawEntries[i]
    24  	}
    25  	m.entries = m.entries[:0]
    26  }
    27  
    28  func (m *OrderedTreeMap[K, V]) Init(cpr Comparator[K], entries int) {
    29  	*m = OrderedTreeMap[K, V]{
    30  		entries: make([]*TreeMapEntry[K, V], entries),
    31  	}
    32  	m.init()
    33  }
    34  
    35  func (m *OrderedTreeMap[K, V]) IsNil() bool {
    36  	return m.entries == nil
    37  }
    38  
    39  func (m *OrderedTreeMap[K, V]) Len() int {
    40  	return len(m.entries)
    41  }
    42  
    43  func (m *OrderedTreeMap[K, V]) Empty() bool {
    44  	return len(m.entries) == 0
    45  }
    46  
    47  func (m *OrderedTreeMap[K, V]) getEntry(key K) *TreeMapEntry[K, V] {
    48  	t := m.root
    49  	if t == nil {
    50  		return nil
    51  	}
    52  	for t != nil {
    53  		if key < t.key {
    54  			t = t.left
    55  		} else if key > t.key {
    56  			t = t.right
    57  		} else {
    58  			return t
    59  		}
    60  	}
    61  	return nil
    62  }
    63  
    64  func (m *OrderedTreeMap[K, V]) Has(key K) bool {
    65  	return m.getEntry(key) != nil
    66  }
    67  
    68  func (m *OrderedTreeMap[K, V]) Get(key K) (value V) {
    69  	if e := m.getEntry(key); e != nil {
    70  		return e.value
    71  	}
    72  	return
    73  }
    74  
    75  func (m *OrderedTreeMap[K, V]) GetEntry(key K) *TreeMapEntry[K, V] {
    76  	return m.getEntry(key)
    77  }
    78  
    79  func (m *OrderedTreeMap[K, V]) Load(key K) (value V, exist bool) {
    80  	if e := m.getEntry(key); e != nil {
    81  		return e.value, true
    82  	}
    83  	return
    84  }
    85  
    86  func (m *OrderedTreeMap[K, V]) Head() (value V, exist bool) {
    87  	if m.head == nil {
    88  		return value, false
    89  	}
    90  	return m.head.value, true
    91  }
    92  
    93  func (m *OrderedTreeMap[K, V]) HeadEntry() *TreeMapEntry[K, V] {
    94  	return m.head
    95  }
    96  
    97  func (m *OrderedTreeMap[K, V]) Tail() (value V, exist bool) {
    98  	if m.tail == nil {
    99  		return value, false
   100  	}
   101  	return m.tail.value, true
   102  }
   103  
   104  func (m *OrderedTreeMap[K, V]) TailEntry() *TreeMapEntry[K, V] {
   105  	return m.tail
   106  }
   107  
   108  func (m *OrderedTreeMap[K, V]) Range(f func(key K, value V) bool) bool {
   109  	var next *TreeMapEntry[K, V]
   110  	for cur := m.head; cur != nil; cur = next {
   111  		next = cur.next()
   112  		if !f(cur.key, cur.value) {
   113  			return false
   114  		}
   115  	}
   116  	return true
   117  }
   118  
   119  func (m *OrderedTreeMap[K, V]) RangeEntries(f func(e *TreeMapEntry[K, V]) bool) bool {
   120  	var next *TreeMapEntry[K, V]
   121  	for cur := m.head; cur != nil; cur = next {
   122  		next = cur.next()
   123  		if !f(cur) {
   124  			return false
   125  		}
   126  	}
   127  	return true
   128  }
   129  
   130  func (m *OrderedTreeMap[K, V]) Keys() (keys []K) {
   131  	var next *TreeMapEntry[K, V]
   132  	for cur := m.head; cur != nil; cur = next {
   133  		next = cur.next()
   134  		keys = append(keys, cur.key)
   135  	}
   136  	return
   137  }
   138  
   139  func (m *OrderedTreeMap[K, V]) Values() (values []V) {
   140  	var next *TreeMapEntry[K, V]
   141  	for cur := m.head; cur != nil; cur = next {
   142  		next = cur.next()
   143  		values = append(values, cur.value)
   144  	}
   145  	return
   146  }
   147  
   148  func (m *OrderedTreeMap[K, V]) grow() int {
   149  	l := len(m.entries)
   150  	if l == cap(m.entries) {
   151  		m.entries = append(m.entries, nil)
   152  		m.entries = m.entries[:cap(m.entries)]
   153  		rawEntries := make([]TreeMapEntry[K, V], cap(m.entries)-l)
   154  		for i := range rawEntries {
   155  			m.entries[i+l] = &rawEntries[i]
   156  		}
   157  		m.entries = m.entries[:l]
   158  	}
   159  	return l
   160  }
   161  
   162  func (m *OrderedTreeMap[K, V]) newEntry(key K, value V, color int) (e *TreeMapEntry[K, V]) {
   163  	l := m.grow()
   164  	m.entries = m.entries[:l+1]
   165  	e = m.entries[l]
   166  	e.key = key
   167  	e.value = value
   168  	e.color = color
   169  	return m.entries[l]
   170  }
   171  
   172  func (m *OrderedTreeMap[K, V]) freeEntry(e *TreeMapEntry[K, V]) {
   173  	l := len(m.entries)
   174  	m.entries[l-1] = e
   175  	m.entries = m.entries[:l-1]
   176  }
   177  
   178  func (*OrderedTreeMap[K, V]) colorOf(p *TreeMapEntry[K, V]) int {
   179  	if p == nil {
   180  		return treeMapColorBlack
   181  	}
   182  	return p.color
   183  }
   184  
   185  func (*OrderedTreeMap[K, V]) parentOf(p *TreeMapEntry[K, V]) *TreeMapEntry[K, V] {
   186  	if p == nil {
   187  		return nil
   188  	}
   189  	return p.parent
   190  }
   191  
   192  func (*OrderedTreeMap[K, V]) setColor(p *TreeMapEntry[K, V], color int) {
   193  	if p != nil {
   194  		p.color = color
   195  	}
   196  }
   197  
   198  func (*OrderedTreeMap[K, V]) leftOf(p *TreeMapEntry[K, V]) *TreeMapEntry[K, V] {
   199  	if p == nil {
   200  		return nil
   201  	}
   202  	return p.left
   203  }
   204  
   205  func (*OrderedTreeMap[K, V]) rightOf(p *TreeMapEntry[K, V]) *TreeMapEntry[K, V] {
   206  	if p == nil {
   207  		return nil
   208  	}
   209  	return p.right
   210  }
   211  
   212  func (m *OrderedTreeMap[K, V]) rotateLeft(p *TreeMapEntry[K, V]) {
   213  	if p != nil {
   214  		r := p.right
   215  		p.right = r.left
   216  		if r.left != nil {
   217  			r.left.parent = p
   218  		}
   219  		r.parent = p.parent
   220  		if p.parent == nil {
   221  			m.root = r
   222  		} else if p.parent.left == p {
   223  			p.parent.left = r
   224  		} else {
   225  			p.parent.right = r
   226  		}
   227  		r.left = p
   228  		p.parent = r
   229  	}
   230  }
   231  
   232  func (m *OrderedTreeMap[K, V]) rotateRight(p *TreeMapEntry[K, V]) {
   233  	if p != nil {
   234  		l := p.left
   235  		p.left = l.right
   236  		if l.right != nil {
   237  			l.right.parent = p
   238  		}
   239  		l.parent = p.parent
   240  		if p.parent == nil {
   241  			m.root = l
   242  		} else if p.parent.right == p {
   243  			p.parent.right = l
   244  		} else {
   245  			p.parent.left = l
   246  		}
   247  		l.right = p
   248  		p.parent = l
   249  	}
   250  }
   251  
   252  func (m *OrderedTreeMap[K, V]) redirectRelation(p, l, r, o, n *TreeMapEntry[K, V]) {
   253  	if p != nil {
   254  		if p.left == o {
   255  			p.left = n
   256  		} else if p.right == o {
   257  			p.right = n
   258  		} else {
   259  			panic("parent not parent of old")
   260  		}
   261  	}
   262  	if l != nil {
   263  		l.parent = n
   264  	}
   265  	if r != nil {
   266  		r.parent = n
   267  	}
   268  }
   269  
   270  func (m *OrderedTreeMap[K, V]) replaceRelation(old, e *TreeMapEntry[K, V]) {
   271  	e.parent = old.parent
   272  	e.left, e.right = old.left, old.right
   273  	e.color = old.color
   274  	m.redirectRelation(old.parent, old.left, old.right, old, e)
   275  	if m.root == old {
   276  		m.root = e
   277  	}
   278  	if m.head == old {
   279  		m.head = e
   280  	}
   281  	if m.tail == old {
   282  		m.tail = e
   283  	}
   284  	old.parent = nil
   285  	old.left, old.right = nil, nil
   286  	old.color = treeMapColorNone
   287  }
   288  
   289  func (m *OrderedTreeMap[K, V]) swapRelation(e1, e2 *TreeMapEntry[K, V]) {
   290  	p1, l1, r1 := e1.parent, e1.left, e1.right
   291  	p2, l2, r2 := e2.parent, e2.left, e2.right
   292  	m.redirectRelation(p1, l1, r1, e1, e2)
   293  	m.redirectRelation(p2, l2, r2, e2, e1)
   294  	e1.parent, e2.parent = e2.parent, e1.parent
   295  	e1.left, e2.left = e2.left, e1.left
   296  	e1.right, e2.right = e2.right, e1.right
   297  	e1.color, e2.color = e2.color, e1.color
   298  	if m.root == e1 {
   299  		m.root = e2
   300  	} else if m.root == e2 {
   301  		m.root = e1
   302  	}
   303  	if m.head == e1 {
   304  		m.head = e2
   305  	} else if m.head == e2 {
   306  		m.head = e1
   307  	}
   308  	if m.tail == e1 {
   309  		m.tail = e2
   310  	} else if m.tail == e2 {
   311  		m.tail = e1
   312  	}
   313  }
   314  
   315  func (m *OrderedTreeMap[K, V]) fixAfterInsertion(x *TreeMapEntry[K, V]) {
   316  	x.color = treeMapColorRed
   317  
   318  	for x != nil && x != m.root && x.parent.color == treeMapColorRed {
   319  		if m.parentOf(x) == m.leftOf(m.parentOf(m.parentOf(x))) {
   320  			y := m.rightOf(m.parentOf(m.parentOf(x)))
   321  			if m.colorOf(y) == treeMapColorRed {
   322  				m.setColor(m.parentOf(x), treeMapColorBlack)
   323  				m.setColor(y, treeMapColorBlack)
   324  				m.setColor(m.parentOf(m.parentOf(x)), treeMapColorRed)
   325  				x = m.parentOf(m.parentOf(x))
   326  			} else {
   327  				if x == m.rightOf(m.parentOf(x)) {
   328  					x = m.parentOf(x)
   329  					m.rotateLeft(x)
   330  				}
   331  				m.setColor(m.parentOf(x), treeMapColorBlack)
   332  				m.setColor(m.parentOf(m.parentOf(x)), treeMapColorRed)
   333  				m.rotateRight(m.parentOf(m.parentOf(x)))
   334  			}
   335  		} else {
   336  			y := m.leftOf(m.parentOf(m.parentOf(x)))
   337  			if m.colorOf(y) == treeMapColorRed {
   338  				m.setColor(m.parentOf(x), treeMapColorBlack)
   339  				m.setColor(y, treeMapColorBlack)
   340  				m.setColor(m.parentOf(m.parentOf(x)), treeMapColorRed)
   341  				x = m.parentOf(m.parentOf(x))
   342  			} else {
   343  				if x == m.leftOf(m.parentOf(x)) {
   344  					x = m.parentOf(x)
   345  					m.rotateRight(x)
   346  				}
   347  				m.setColor(m.parentOf(x), treeMapColorBlack)
   348  				m.setColor(m.parentOf(m.parentOf(x)), treeMapColorRed)
   349  				m.rotateLeft(m.parentOf(m.parentOf(x)))
   350  			}
   351  		}
   352  	}
   353  	m.root.color = treeMapColorBlack
   354  }
   355  
   356  func (m *OrderedTreeMap[K, V]) addEntryToEmpty(e *TreeMapEntry[K, V]) {
   357  	m.root = e
   358  	m.head, m.tail = e, e
   359  }
   360  
   361  func (m *OrderedTreeMap[K, V]) addEntryToParent(e, p *TreeMapEntry[K, V], addToLeft bool) {
   362  	e.parent = p
   363  	if addToLeft {
   364  		p.left = e
   365  		if p == m.head {
   366  			m.head = e
   367  		}
   368  	} else {
   369  		p.right = e
   370  		if p == m.tail {
   371  			m.tail = e
   372  		}
   373  	}
   374  	m.fixAfterInsertion(e)
   375  }
   376  
   377  func (m *OrderedTreeMap[K, V]) add(key K, value V, mode int) (old V, exist bool) {
   378  	t := m.root
   379  	if t == nil {
   380  		m.addEntryToEmpty(m.newEntry(key, value, treeMapColorBlack))
   381  		return value, false
   382  	}
   383  	var parent *TreeMapEntry[K, V]
   384  	var cmp int
   385  	for t != nil {
   386  		parent = t
   387  		if key < t.key {
   388  			t = t.left
   389  			cmp = -1
   390  		} else if key > t.key {
   391  			t = t.right
   392  			cmp = 1
   393  		} else {
   394  			switch mode {
   395  			case treeMapPut:
   396  				m.replaceRelation(t, m.newEntry(key, value, treeMapColorBlack))
   397  				return value, false
   398  			case treeMapReplace, treeMapStore:
   399  				old = t.value
   400  				t.value = value
   401  				return old, true
   402  			case treeMapLoadOrStore:
   403  				return t.value, true
   404  			}
   405  		}
   406  	}
   407  	if mode != treeMapReplace {
   408  		m.addEntryToParent(m.newEntry(key, value, treeMapColorBlack), parent, cmp < 0)
   409  		return value, false
   410  	}
   411  	return
   412  }
   413  
   414  func (m *OrderedTreeMap[K, V]) Store(key K, value V) {
   415  	m.add(key, value, treeMapStore)
   416  }
   417  
   418  func (m *OrderedTreeMap[K, V]) Set(key K, value V) (old V, exist bool) {
   419  	return m.add(key, value, treeMapStore)
   420  }
   421  
   422  func (m *OrderedTreeMap[K, V]) LoadOrStore(key K, value V) (old V, exist bool) {
   423  	return m.add(key, value, treeMapLoadOrStore)
   424  }
   425  
   426  func (m *OrderedTreeMap[K, V]) Replace(ket K, value V) bool {
   427  	_, exist := m.add(ket, value, treeMapReplace)
   428  	return !exist
   429  }
   430  
   431  func (m *OrderedTreeMap[K, V]) Put(key K, value V) (old V, exist bool) {
   432  	return m.add(key, value, treeMapPut)
   433  }
   434  
   435  func (m *OrderedTreeMap[K, V]) fixAfterDeletion(x *TreeMapEntry[K, V]) {
   436  	for x != m.root && m.colorOf(x) == treeMapColorBlack {
   437  		if x == m.leftOf(m.parentOf(x)) {
   438  			sib := m.rightOf(m.parentOf(x))
   439  
   440  			if m.colorOf(sib) == treeMapColorRed {
   441  				m.setColor(sib, treeMapColorBlack)
   442  				m.setColor(m.parentOf(x), treeMapColorRed)
   443  				m.rotateLeft(m.parentOf(x))
   444  				sib = m.rightOf(m.parentOf(x))
   445  			}
   446  
   447  			if m.colorOf(m.leftOf(sib)) == treeMapColorBlack &&
   448  				m.colorOf(m.rightOf(sib)) == treeMapColorBlack {
   449  				m.setColor(sib, treeMapColorRed)
   450  				x = m.parentOf(x)
   451  			} else {
   452  				if m.colorOf(m.rightOf(sib)) == treeMapColorBlack {
   453  					m.setColor(m.leftOf(sib), treeMapColorBlack)
   454  					m.setColor(sib, treeMapColorRed)
   455  					m.rotateRight(sib)
   456  					sib = m.rightOf(m.parentOf(x))
   457  				}
   458  				m.setColor(sib, m.colorOf(m.parentOf(x)))
   459  				m.setColor(m.parentOf(x), treeMapColorBlack)
   460  				m.setColor(m.rightOf(sib), treeMapColorBlack)
   461  				m.rotateLeft(m.parentOf(x))
   462  				x = m.root
   463  			}
   464  		} else { // symmetric
   465  			sib := m.leftOf(m.parentOf(x))
   466  
   467  			if m.colorOf(sib) == treeMapColorRed {
   468  				m.setColor(sib, treeMapColorBlack)
   469  				m.setColor(m.parentOf(x), treeMapColorRed)
   470  				m.rotateRight(m.parentOf(x))
   471  				sib = m.leftOf(m.parentOf(x))
   472  			}
   473  
   474  			if m.colorOf(m.rightOf(sib)) == treeMapColorBlack &&
   475  				m.colorOf(m.leftOf(sib)) == treeMapColorBlack {
   476  				m.setColor(sib, treeMapColorRed)
   477  				x = m.parentOf(x)
   478  			} else {
   479  				if m.colorOf(m.leftOf(sib)) == treeMapColorBlack {
   480  					m.setColor(m.rightOf(sib), treeMapColorBlack)
   481  					m.setColor(sib, treeMapColorRed)
   482  					m.rotateLeft(sib)
   483  					sib = m.leftOf(m.parentOf(x))
   484  				}
   485  				m.setColor(sib, m.colorOf(m.parentOf(x)))
   486  				m.setColor(m.parentOf(x), treeMapColorBlack)
   487  				m.setColor(m.leftOf(sib), treeMapColorBlack)
   488  				m.rotateRight(m.parentOf(x))
   489  				x = m.root
   490  			}
   491  		}
   492  	}
   493  
   494  	m.setColor(x, treeMapColorBlack)
   495  }
   496  
   497  func (m *OrderedTreeMap[K, V]) deleteEntry(p *TreeMapEntry[K, V]) {
   498  	// If strictly internal, copy successor's element to p and then make p
   499  	// point to successor.
   500  	if p.left != nil && p.right != nil {
   501  		m.swapRelation(p, p.next())
   502  	} // p has 2 children
   503  
   504  	if p == m.head {
   505  		m.head = p.next()
   506  	}
   507  	if p == m.tail {
   508  		m.tail = p.prev()
   509  	}
   510  
   511  	// Start fixup at replacement node, if it exists.
   512  	var replacement *TreeMapEntry[K, V]
   513  	if p.left != nil {
   514  		replacement = p.left
   515  	} else {
   516  		replacement = p.right
   517  	}
   518  
   519  	if replacement != nil {
   520  		// Link replacement to parent
   521  		replacement.parent = p.parent
   522  		if p.parent == nil {
   523  			m.root = replacement
   524  		} else if p == p.parent.left {
   525  			p.parent.left = replacement
   526  		} else {
   527  			p.parent.right = replacement
   528  		}
   529  
   530  		// Null out links so they are OK to use by fixAfterDeletion.
   531  		p.left, p.right, p.parent = nil, nil, nil
   532  
   533  		// Fix replacement
   534  		if p.color == treeMapColorBlack {
   535  			m.fixAfterDeletion(replacement)
   536  		}
   537  	} else if p.parent == nil { // return if we are the only node.
   538  		m.root = nil
   539  	} else { //  No children. Use self as phantom replacement and unlink.
   540  		if p.color == treeMapColorBlack {
   541  			m.fixAfterDeletion(p)
   542  		}
   543  
   544  		if p.parent != nil {
   545  			if p == p.parent.left {
   546  				p.parent.left = nil
   547  			} else if p == p.parent.right {
   548  				p.parent.right = nil
   549  			}
   550  			p.parent = nil
   551  		}
   552  	}
   553  
   554  	p.color = treeMapColorNone
   555  	m.freeEntry(p)
   556  }
   557  
   558  func (m *OrderedTreeMap[K, V]) Remove(key K) (old V, exist bool) {
   559  	if e := m.getEntry(key); e != nil {
   560  		old, exist = e.value, true
   561  		m.deleteEntry(e)
   562  	}
   563  	return
   564  }
   565  
   566  func (m *OrderedTreeMap[K, V]) RemoveEntry(e *TreeMapEntry[K, V]) bool {
   567  	if e.Removed() {
   568  		return false
   569  	}
   570  	m.deleteEntry(e)
   571  	return true
   572  }
   573  
   574  func (m *OrderedTreeMap[K, V]) Delete(key K) {
   575  	if e := m.getEntry(key); e != nil {
   576  		m.deleteEntry(e)
   577  	}
   578  	return
   579  }
   580  
   581  func (m *OrderedTreeMap[K, V]) Shift() (v V, exist bool) {
   582  	if m.head != nil {
   583  		v, exist = m.head.value, true
   584  		m.deleteEntry(m.head)
   585  	}
   586  	return
   587  }
   588  
   589  func (m *OrderedTreeMap[K, V]) ShiftKey() (k K, exist bool) {
   590  	if m.head != nil {
   591  		k, exist = m.head.key, true
   592  		m.deleteEntry(m.head)
   593  	}
   594  	return
   595  }
   596  
   597  func (m *OrderedTreeMap[K, V]) Pop() (v V, exist bool) {
   598  	if m.tail != nil {
   599  		v, exist = m.tail.value, true
   600  		m.deleteEntry(m.tail)
   601  	}
   602  	return
   603  }
   604  
   605  func (m *OrderedTreeMap[K, V]) PopKey() (k K, exist bool) {
   606  	if m.tail != nil {
   607  		k, exist = m.tail.key, true
   608  		m.deleteEntry(m.tail)
   609  	}
   610  	return
   611  }
   612  
   613  func (m *OrderedTreeMap[K, V]) Clear() {
   614  	for _, entry := range m.entries {
   615  		entry.color = treeMapColorNone
   616  		entry.parent, entry.left, entry.right = nil, nil, nil
   617  	}
   618  	m.entries = m.entries[:0]
   619  	m.root, m.head, m.tail = nil, nil, nil
   620  }