Golang优先队列与有序集合的实现

2 minute read

Published:

Golang优先队列与有序集合的实现

题目链接:https://leetcode.cn/problems/stock-price-fluctuation/description/

优先队列

package main

import "container/heap"

type StockPrice struct {
	timePriceMap map[int]int
	maxPrice     hp
	minPrice     hp
	latest       int
}

func Constructor() StockPrice {
	st := StockPrice{}
	st.timePriceMap = make(map[int]int)
	st.latest = 0
	return st
}

func (this *StockPrice) Update(timestamp int, price int) {
	this.latest = max(this.latest, timestamp)
	this.timePriceMap[timestamp] = price
	heap.Push(&this.maxPrice, pair{-price, timestamp})
	heap.Push(&this.minPrice, pair{price, timestamp})
}

func (this *StockPrice) Current() int {
	return this.timePriceMap[this.latest]
}

func (this *StockPrice) Maximum() int {
	for this.maxPrice.Len() > 0 {
		if top := this.maxPrice[0]; -top.price == this.timePriceMap[top.timestamp] {
			return -top.price
		}
		heap.Pop(&this.maxPrice)
	}
	return -1
}

func (this *StockPrice) Minimum() int {
	for this.minPrice.Len() > 0 {
		if top := this.minPrice[0]; top.price == this.timePriceMap[top.timestamp] {
			return top.price
		}
		heap.Pop(&this.minPrice)
	}
	return -1
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

type pair struct{ price, timestamp int }
type hp []pair

// 递增优先队列
func (h hp) Len() int            { return len(h) }
func (h hp) Less(i, j int) bool  { return h[i].price < h[j].price }
func (h hp) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *hp) Push(v interface{}) { *h = append(*h, v.(pair)) }
func (h *hp) Pop() interface{}   { a := *h; v := a[len(a)-1]; *h = a[:len(a)-1]; return v }

此方法实现为递增优先队列

红黑树


package main

import "github.com/emirpasic/gods/trees/redblacktree"

type StockPrice struct {
	prices       *redblacktree.Tree // {key, value}形式, 按照key从小到大排序
	timePriceMap map[int]int
	maxTimeStamp int
}

func Constructor() StockPrice {
	return StockPrice{redblacktree.NewWithIntComparator(), map[int]int{}, 0}
}

func (this *StockPrice) Update(timestamp int, price int) {
	this.maxTimeStamp = max(this.maxTimeStamp, timestamp)
	if prevPrice := this.timePriceMap[timestamp]; prevPrice > 0 { // 之前存在, 需要将之前价格的计数器-1
		if times, _ := this.prices.Get(prevPrice); times.(int) > 1 { // 获取当前价格出现的次数
			// 将times-1
			this.prices.Put(prevPrice, times.(int)-1)
		} else {
			// times == 1, 将该键值对移除
			this.prices.Remove(prevPrice)
		}
	}
	count := 0 // 计算当前price的出现在prices中次数
	if val, ok := this.prices.Get(price); ok {
		count = val.(int)
	}
	this.prices.Put(price, count+1)
	this.timePriceMap[timestamp] = price

}

func (this *StockPrice) Current() int {
	return this.timePriceMap[this.maxTimeStamp]
}

func (this *StockPrice) Maximum() int {
	return this.prices.Right().Key.(int)
}

func (this *StockPrice) Minimum() int {
	return this.prices.Left().Key.(int)
}

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}

*redblacktree.Tree为键值有序的hashMap, 可用于实现有序集合