用go实现一个LFU缓存

LFU(Least Frenquently Used)，和 LRU 类似，也是一种常用的缓存机制和算法，在数据数量达到缓存容量时最少使用的会优先淘汰。

简介

其形式如下：

实现设计

为了高效率的查询和写入，可以像 LRU 那样使用哈希表。而 LFU 需要把最少使用的优先淘汰，那么数据需要记录使用次数，再者当有两个相同使用次数时怎么办呢？再比较下添加时间好了，早添加的优先被移除。为了方便移除需要对数据排序，需要选择一种排序高效的数据结构，于是可以使用优先队列（通常最小二叉堆或最大二叉堆实现），排序的平均时间复杂度为 \(O(nlogn)\)。

实现

先定义数据元素：

type lfuElement struct {
	key    string
	value  interface{}
	aCount int   // 访问数量
	aTime  int64 // 添加时间戳
	idx    int	 // 队列中的位置
}

优先队列

type priorityQueue []*lfuElement

// 队列长度
func (q priorityQueue) Len() int {
	return len(q)
}

// Less 与 Swap 定义了排序的逻辑

func (q priorityQueue) Less(i, j int) bool {
	if q[i].aCount == q[j].aCount {
		return q[i].aTime < q[j].aTime
	}
	return q[i].aCount < q[j].aCount
}

// 交换元素
func (q priorityQueue) Swap(i, j int) {
	// 交换元素
	q[i], q[j] = q[j], q[i]
	// 索引不用交换
	q[i].idx = i
	q[j].idx = j
}

// 添加数据
func (q *priorityQueue) Push(v interface{}) {
	l := q.Len()
	en := v.(*lfuElement)
	en.idx = l
	*q = append(*q, en) // 这里会重新分配内存，并拷贝数据
}

// 弹出数据
func (q *priorityQueue) Pop() interface{} {
	tmpQ := *q
	l := len(tmpQ)
	en := tmpQ[l-1]
	tmpQ[l-1] = nil
	*q = tmpQ[:l-1]
	return en
}

其实上面的 priorityQueue 类型实现了 heap 的接口：

// The Interface type describes the requirements
// for a type using the routines in this package.
// Any type that implements it may be used as a
// min-heap with the following invariants (established after
// Init has been called or if the data is empty or sorted):
//
//	!h.Less(j, i) for 0 <= i < h.Len() and 2*i+1 <= j <= 2*i+2 and j < h.Len()
//
// Note that Push and Pop in this interface are for package heap's
// implementation to call. To add and remove things from the heap,
// use heap.Push and heap.Pop.
type Interface interface {
	sort.Interface
	Push(x interface{}) // add x as element Len()
	Pop() interface{}   // remove and return element Len() - 1.
}

定义LFU缓存：

type LFUCache struct {
	dataMap   map[string]*lfuElement
	dataQueue *priorityQueue
	capacity  int
}

// 初始化
func (cache *LFUCache) Init(capacity int) {
	cache.capacity = capacity
	cache.dataMap = make(map[string]*lfuElement)
	q := make(priorityQueue, 0)
	cache.dataQueue = &q
}

// 添加/更新数据
func (cache *LFUCache) Put(key string, value interface{}) {
	if ele, ok := cache.dataMap[key]; ok {
		// key 存在，更新
		ele.value = value
		ele.aCount++
		// 重排序
		heap.Fix(cache.dataQueue, ele.idx)
	} else {
		// 不存在，添加
		ele := &lfuElement{
			key:    key,
			value:  value,
			aCount: 1,
			aTime:  time.Now().UnixNano(),
		}
		cache.dataQueue.Push(ele)
		cache.dataMap[key] = ele
		if cache.dataQueue.Len() > cache.capacity {
			// 超出了容量限制，移除最少使用的
			cache.rmLeastFrequent()
		}
	}
}

// 读取数据
func (cache *LFUCache) Get(key string) interface{} {
	if ele, ok := cache.dataMap[key]; ok {
		ele.aCount++ // 使用计数+1
		// 重排序
		heap.Fix(cache.dataQueue, ele.idx)
		return ele.value
	}
	return nil
}

// 获取数据容量
func (cache *LFUCache) Capacity() int {
	return cache.capacity
}

// 获取数据长度
func (cache *LFUCache) Length() int {
	return cache.dataQueue.Len()
}

// 获取所有数据，顺序展示
func (cache *LFUCache) ToList() []interface{} {
	result := make([]interface{}, cache.dataQueue.Len())
	tmpQ := *cache.dataQueue
	sort.Sort(tmpQ)
	for idx, ele := range tmpQ {
		result[idx] = ele.value
	}
	return result
}

// 删除最少使用的
func (cache *LFUCache) rmLeastFrequent() {
	if cache.dataQueue.Len() == 0 {
		// 空的，不处理
		return
	}
	// 从queue中删除
	val := heap.Pop(cache.dataQueue)
	ele := val.(*lfuElement)
	// 从map中删除
	delete(cache.dataMap, ele.key)
}

测试：

func main() {
	var lfu LFUCache
	lfu.Init(5)
	lfu.Put("a", "apple")
	fmt.Printf("lfu set a :%v \n", "apple")
	lfu.Put("b", "big")
	fmt.Printf("lfu set b :%v \n", "big")
	lfu.Put("c", "car")
	fmt.Printf("lfu set c :%v \n", "car")
	lfu.Put("d", "do")
	fmt.Printf("lfu set d :%v \n", "do")
	lfu.Put("e", "eye")
	fmt.Printf("lfu set e :%v \n", "eye")
	lfu.Put("f", "flu") // 增加f，a 应该被淘汰
	fmt.Printf("lfu set f :%v \n", "flu")

	fmt.Printf("lfu content:%v \n", lfu.ToList())

	fmt.Printf("lfu get b :%v \n", lfu.Get("b")) // b的访问次数+1

	fmt.Printf("lfu content:%v \n", lfu.ToList())
	lfu.Put("b", "book") // b的访问次数+1，并更新值
	fmt.Printf("lfu set b :%v \n", "book")
	fmt.Printf("lfu content:%v \n", lfu.ToList())

	fmt.Printf("lfu get b :%v \n", lfu.Get("b")) // b的访问次数+1
	fmt.Printf("lfu get c :%v \n", lfu.Get("c")) // c的访问次数+1
	fmt.Printf("lfu get d :%v \n", lfu.Get("d")) // d的访问次数+1
	fmt.Printf("lfu get d :%v \n", lfu.Get("d")) // d的访问次数+1
	fmt.Printf("lfu content:%v \n", lfu.ToList())
	lfu.Put("g", "gas") // 增加 g，淘汰e
	fmt.Printf("lfu set g :%v \n", "gas")
	fmt.Printf("lfu content:%v \n", lfu.ToList())
}

输出：

lfu set a :apple
lfu set b :big
lfu set c :car
lfu set d :do
lfu set e :eye
lfu set f :flu
lfu content:[big car do eye flu]
lfu get b :big
lfu content:[car do eye flu big]
lfu set b :book
lfu content:[car do eye flu book]
lfu get b :book
lfu get c :car
lfu get d :do
lfu get d :do
lfu content:[eye flu car do book]
lfu set g :gas
lfu content:[flu gas car do book]

这样一个基本的LFU缓存就做好了。