Probabilistic structures: filters, sketches and samples

// BloomFor sizes a filter for n items at false positive rate p:
// m = -n ln p / (ln 2)^2 bits, and k = (m/n) ln 2 positions per item.
func BloomFor(n int, p float64) *Bloom {
  m := uint64(math.Ceil(-float64(n) * math.Log(p) / (math.Ln2 * math.Ln2)))
  return NewBloom(m, OptimalK(float64(m)/float64(n)))
}

// OptimalK is the k that minimises the false positive rate at b bits per item: b ln 2, rounded.
func OptimalK(b float64) int {
  return max(1, int(math.Round(b*math.Ln2)))
}

// FalsePositiveRate is the chance that an item never added finds all its k bits set, after n
// items went into m bits: (1 - e^(-kn/m))^k.
func FalsePositiveRate(m uint64, n int, k int) float64 {
  return math.Pow(1-math.Exp(-float64(k)*float64(n)/float64(m)), float64(k))
}

// Positions returns the k bit positions of a hashed item: h1 + i*h2 mod m (double hashing).
func (b *Bloom) Positions(h uint64) []uint64 {
  h1, h2 := pair(h)
  out := make([]uint64, b.k)
  for i := range b.k {
    out[i] = (h1 + uint64(i)*h2) % b.m
  }
  return out
}

// AddHash sets the k bits of a hashed item.
func (b *Bloom) AddHash(h uint64) {
  h1, h2 := pair(h)
  for i := range b.k {
    p := (h1 + uint64(i)*h2) % b.m
    b.bits[p/64] |= 1 << (p % 64)
  }
  b.n++
}

// HasHash reports whether all k bits of a hashed item are set. It stops at the first clear bit,
// so a definite "no" usually costs one or two reads.
func (b *Bloom) HasHash(h uint64) bool {
  h1, h2 := pair(h)
  for i := range b.k {
    p := (h1 + uint64(i)*h2) % b.m
    if b.bits[p/64]&(1<<(p%64)) == 0 {
      return false
    }
  }
  return true
}

-- Add an item to a Bloom filter kept in a Redis bitmap, and say whether it was new.
-- KEYS[1] the bitmap. ARGV holds the item's k bit positions, computed by the client.
-- Returns the number of bits that were 0 before. 0 means every bit was set: the item was
-- probably seen before. Redis runs the script as one step, so two callers that add the
-- same new item get one 'new' answer between them.
local new = 0
for i = 1, #ARGV do
  if redis.call('SETBIT', KEYS[1], ARGV[i], 1) == 0 then
    new = new + 1
  end
end
return new

-- Ask a Bloom filter kept in a Redis bitmap whether it may hold an item.
-- KEYS[1] the bitmap. ARGV holds the item's k bit positions, computed by the client.
-- Returns 0 at the first clear bit (never added) and 1 if every bit is set (maybe added).
for i = 1, #ARGV do
  if redis.call('GETBIT', KEYS[1], ARGV[i]) == 0 then
    return 0
  end
end
return 1

// AddHash routes the hash to register j (its first p bits) and keeps the largest rank seen there.
// The rank is the position of the first 1 bit in the remaining 64 - p bits.
func (h *HLL) AddHash(x uint64) {
  j := x >> (64 - h.p)
  rest := x<<h.p | 1<<(h.p-1) // the guard bit caps the rank at 64 - p + 1
  rank := uint8(bits.LeadingZeros64(rest)) + 1
  if rank > h.reg[j] {
    h.reg[j] = rank
  }
}

// Count is the harmonic-mean estimate α·m² / Σ 2^-reg. While many registers are still 0, it uses
// linear counting, m·ln(m / zeros), which is more accurate for small sets.
func (h *HLL) Count() uint64 {
  m := float64(len(h.reg))
  sum, zeros := 0.0, 0
  for _, r := range h.reg {
    sum += math.Ldexp(1, -int(r))
    if r == 0 {
      zeros++
    }
  }
  est := alpha(len(h.reg)) * m * m / sum
  if est <= 2.5*m && zeros > 0 {
    est = m * math.Log(m/float64(zeros))
  }
  return uint64(math.Round(est))
}

// Merge keeps the larger value of each register, so the result counts the union of both sets.
func (h *HLL) Merge(o *HLL) error {
  if o.p != h.p {
    return fmt.Errorf("merge hll: p %d into p %d", o.p, h.p)
  }
  for i, r := range o.reg {
    h.reg[i] = max(h.reg[i], r)
  }
  return nil
}

// AddHash adds c to one counter in each row.
func (s *CountMin) AddHash(h uint64, c uint32) {
  h1, h2 := pair(h)
  for i := range s.d {
    s.rows[i][(h1+uint64(i)*h2)%uint64(s.w)] += c
  }
  s.total += uint64(c)
}

// EstimateHash is the smallest of the item's d counters: never below the true count.
func (s *CountMin) EstimateHash(h uint64) uint32 {
  h1, h2 := pair(h)
  est := uint32(math.MaxUint32)
  for i := range s.d {
    est = min(est, s.rows[i][(h1+uint64(i)*h2)%uint64(s.w)])
  }
  return est
}

// index returns the item's fingerprint (never 0, since 0 marks an empty slot) and both buckets.
func (c *Cuckoo) index(h uint64) (fp uint16, i1, i2 uint64) {
  fp = uint16(h>>(64-c.fpBits)) & (1<<c.fpBits - 1)
  if fp == 0 {
    fp = 1
  }
  i1 = h & c.mask
  return fp, i1, c.alt(i1, fp)
}

// alt is the other bucket of a fingerprint in bucket i. alt(alt(i, fp), fp) == i.
func (c *Cuckoo) alt(i uint64, fp uint16) uint64 {
  return (i ^ mix(uint64(fp))) & c.mask
}

// Add stores the fingerprint of s in a free slot of either bucket. When both are full, it evicts
// a random fingerprint, moves it to its other bucket, and repeats up to maxKicks times.
func (c *Cuckoo) Add(s string) error {
  fp, i1, i2 := c.index(Hash(s))
  if c.put(i1, fp) || c.put(i2, fp) {
    c.n++
    return nil
  }
  i := i1
  if c.rnd.IntN(2) == 1 {
    i = i2
  }
  for range c.maxKicks {
    slot := c.rnd.IntN(4)
    fp, c.buckets[i][slot] = c.buckets[i][slot], fp
    i = c.alt(i, fp)
    if c.put(i, fp) {
      c.n++
      return nil
    }
  }
  return ErrFull // fp, the last evicted fingerprint, is lost; the filter now has a false negative
}

// Has reports whether either bucket of s holds its fingerprint.
func (c *Cuckoo) Has(s string) bool {
  fp, i1, i2 := c.index(Hash(s))
  return c.holds(i1, fp) || c.holds(i2, fp)
}

// Remove deletes one copy of the fingerprint of s. Remove only items that were added.
func (c *Cuckoo) Remove(s string) bool {
  fp, i1, i2 := c.index(Hash(s))
  for _, i := range []uint64{i1, i2} {
    for j, v := range c.buckets[i] {
      if v == fp {
        c.buckets[i][j] = 0
        c.n--
        return true
      }
    }
  }
  return false
}

// Add counts one occurrence of key. A key in the table gets +1. A new key takes a free counter,
// or replaces the key with the smallest count and inherits that count + 1. The inherited part is
// its error bound.
func (s *SpaceSaving) Add(key string) {
  s.n++
  if i, ok := s.h.at[key]; ok {
    s.h.items[i].Count++
    heap.Fix(&s.h, i)
    return
  }
  if s.h.Len() < s.cap {
    heap.Push(&s.h, &Item{Key: key, Count: 1})
    return
  }
  low := s.h.items[0]
  delete(s.h.at, low.Key)
  s.h.at[key] = 0
  *low = Item{Key: key, Count: low.Count + 1, Err: low.Count}
  heap.Fix(&s.h, 0)
}

// Add counts key in the sketch, then offers its new estimate to the heap. The heap admits a key
// whose estimate beats the smallest of the current k.
func (t *CMSTopK) Add(key string) {
  h := Hash(key)
  t.cms.AddHash(h, 1)
  est := uint64(t.cms.EstimateHash(h))
  if i, ok := t.h.at[key]; ok {
    t.h.items[i].Count = est
    heap.Fix(&t.h, i)
    return
  }
  if t.h.Len() < t.k {
    heap.Push(&t.h, &Item{Key: key, Count: est})
    return
  }
  if est > t.h.items[0].Count {
    heap.Pop(&t.h)
    heap.Push(&t.h, &Item{Key: key, Count: est})
  }
}

// Add offers item number n (counting from 1). The first k items fill the sample. After that, item
// n replaces a random slot with probability k/n.
func (r *Reservoir[T]) Add(item T) {
  r.seen++
  if len(r.out) < r.k {
    r.out = append(r.out, item)
    return
  }
  if j := r.rnd.IntN(r.seen); j < r.k {
    r.out[j] = item
  }
}