cache_lfu.go

package main

import "fmt"

type Set[T comparable] struct {
	items map[T]struct{}
}

func NewSet[T comparable]() *Set[T] {
	return &Set[T]{
		items: make(map[T]struct{}),
	}
}

func (s *Set[T]) Add(val T) {
	s.items[val] = struct{}{}
}

func (s *Set[T]) Remove(val T) {
	delete(s.items, val)
}

func (s *Set[T]) Pop() T {
	if len(s.items) == 0 {
		panic("Pop() from the emty Set")
	}
	for key := range s.items {
		s.Remove(key)
		return key
	}
	var ans T
	return ans
}

func (s *Set[T]) IsEmpty() bool {
	return len(s.items) == 0
}

func (s *Set[T]) String() string {
	return fmt.Sprintf("%v", s.items)
}

type LFUCache[KeyT comparable, ValueT any] struct {
	capacity int
	values   map[KeyT]ValueT
	freqs    map[KeyT]int
	keys     map[int]*Set[KeyT]
	minFreq  int
}

func NewLFUCache[KeyT comparable, ValueT any](capacity int) *LFUCache[KeyT, ValueT] {
	return &LFUCache[KeyT, ValueT]{
		capacity: capacity,
		values:   make(map[KeyT]ValueT),
		freqs:    make(map[KeyT]int),
		keys:     make(map[int]*Set[KeyT]),
		minFreq:  0,
	}
}

func (lfu *LFUCache[KeyT, ValueT]) Get(key KeyT) (any, bool) {
	if _, ok := lfu.values[key]; !ok {
		return nil, false
	}
	value := lfu.values[key]
	freq := lfu.freqs[key]
	lfu.update(key, value, freq)
	return value, true
}

func (lfu *LFUCache[KeyT, ValueT]) Put(key KeyT, value ValueT) {
	if _, ok := lfu.freqs[key]; ok {
		freq := lfu.freqs[key]
		lfu.update(key, value, freq)
	} else {
		if len(lfu.values) >= lfu.capacity {
			lfu.evict()
		}
		lfu.values[key] = value
		lfu.freqs[key] = 1
		if _, ok := lfu.keys[1]; !ok {
			lfu.keys[1] = NewSet[KeyT]()
		}
		lfu.keys[1].Add(key)
		lfu.minFreq = 1
	}
}

func (lfu *LFUCache[KeyT, ValueT]) update(key KeyT, value ValueT, freq int) {
	lfu.keys[freq].Remove(key)
	if lfu.keys[freq].IsEmpty() {
		delete(lfu.keys, freq)
		if lfu.minFreq == freq {
			lfu.minFreq += 1
		}
	}

	lfu.values[key] = value
	lfu.freqs[key] = freq + 1
	if _, ok := lfu.keys[freq+1]; !ok {
		lfu.keys[freq+1] = NewSet[KeyT]()
	}
	lfu.keys[freq+1].Add(key)
}

func (lfu *LFUCache[KeyT, ValueT]) evict() {
	key := lfu.keys[lfu.minFreq].Pop()
	if lfu.keys[lfu.minFreq].IsEmpty() {
		delete(lfu.keys, lfu.minFreq)
	}
	delete(lfu.values, key)
	delete(lfu.freqs, key)
}

func main() {
	Assert := func(ok bool) {
		if !ok {
			panic("Assertion is failed")
		}
	}

	lfu := NewLFUCache[int, string](3)

	lfu.Put(1, "A")
	lfu.Put(2, "B")

	got, ok := lfu.Get(1)
	Assert(ok && got == "A")

	lfu.Put(3, "C")

	got, ok = lfu.Get(2)
	Assert(ok && got == "B")

	lfu.Put(4, "D")

	got, ok = lfu.Get(3)
	Assert(!ok)

	got, ok = lfu.Get(1)
	Assert(ok && got == "A")

	got, ok = lfu.Get(2)
	Assert(ok && got == "B")

	got, ok = lfu.Get(4)
	Assert(ok && got == "D")

	fmt.Println("OK")
}