Locality-Sensitive Hashing (LSH) — Interview Preparation¶

package main

import (
    "fmt"
    "math"
    "math/rand"
)

type LSH struct {
    b, r    int
    planes  [][]float64        // (b*r) random normals
    tables  []map[uint64][]int // one per band
    vecs    [][]float64
    ids     []int
}

func NewLSH(b, r, d int, seed int64) *LSH {
    rng := rand.New(rand.NewSource(seed))
    k := b * r
    planes := make([][]float64, k)
    for i := range planes {
        planes[i] = make([]float64, d)
        for j := range planes[i] {
            planes[i][j] = rng.NormFloat64()
        }
    }
    t := make([]map[uint64][]int, b)
    for i := range t {
        t[i] = map[uint64][]int{}
    }
    return &LSH{b: b, r: r, planes: planes, tables: t}
}

func (l *LSH) bandKey(x []float64, band int) uint64 {
    var key uint64
    for row := 0; row < l.r; row++ {
        p := l.planes[band*l.r+row]
        dot := 0.0
        for j := range x {
            dot += x[j] * p[j]
        }
        key <<= 1
        if dot >= 0 {
            key |= 1
        }
    }
    return key | (uint64(band) << 40) // namespace per band
}

func (l *LSH) Add(id int, x []float64) {
    idx := len(l.vecs)
    l.vecs = append(l.vecs, x)
    l.ids = append(l.ids, id)
    for band := 0; band < l.b; band++ {
        k := l.bandKey(x, band)
        l.tables[band][k] = append(l.tables[band][k], idx)
    }
}

func cosine(a, b []float64) float64 {
    var dot, na, nb float64
    for i := range a {
        dot += a[i] * b[i]
        na += a[i] * a[i]
        nb += b[i] * b[i]
    }
    return dot / (math.Sqrt(na) * math.Sqrt(nb))
}

func (l *LSH) Query(x []float64) (int, float64) {
    seen := map[int]bool{}
    best, bestSim := -1, -2.0
    for band := 0; band < l.b; band++ {
        k := l.bandKey(x, band)
        for _, idx := range l.tables[band][k] {
            if seen[idx] {
                continue
            }
            seen[idx] = true
            s := cosine(x, l.vecs[idx])
            if s > bestSim {
                bestSim, best = s, l.ids[idx]
            }
        }
    }
    return best, bestSim
}

func main() {
    l := NewLSH(8, 4, 5, 42) // k = 32 bits
    l.Add(1, []float64{1, 0, 0, 0, 0})
    l.Add(2, []float64{0.98, 0.05, 0.0, 0.0, 0.1}) // near-dup of id 1
    l.Add(3, []float64{0, 0, 1, 1, 0})             // unrelated
    id, sim := l.Query([]float64{0.97, 0.04, 0, 0, 0.08})
    fmt.Printf("best id=%d cosine=%.4f\n", id, sim) // expect id 2 (or 1), high cosine
}

package main

import (
    "fmt"
    "hash/fnv"
)

type MinHashLSH struct {
    b, r  int
    seeds []uint64
    bands map[string][]int // band-bucket -> ids
}

func NewMinHashLSH(b, r int) *MinHashLSH {
    k := b * r
    seeds := make([]uint64, k)
    for i := range seeds {
        seeds[i] = uint64(i)*2654435761 + 1 // distinct odd multipliers
    }
    return &MinHashLSH{b: b, r: r, seeds: seeds, bands: map[string][]int{}}
}

func hashShingle(s uint64, seed uint64) uint64 {
    h := fnv.New64a()
    var buf [16]byte
    for i := 0; i < 8; i++ {
        buf[i] = byte(s >> (8 * i))
        buf[8+i] = byte(seed >> (8 * i))
    }
    h.Write(buf[:])
    return h.Sum64()
}

func (m *MinHashLSH) signature(set []uint64) []uint64 {
    k := m.b * m.r
    sig := make([]uint64, k)
    for i := range sig {
        sig[i] = ^uint64(0) // +inf
    }
    for _, sh := range set {
        for i := 0; i < k; i++ {
            if h := hashShingle(sh, m.seeds[i]); h < sig[i] {
                sig[i] = h
            }
        }
    }
    return sig
}

func (m *MinHashLSH) Add(id int, set []uint64) {
    sig := m.signature(set)
    for band := 0; band < m.b; band++ {
        key := fmt.Sprintf("%d:", band)
        for row := 0; row < m.r; row++ {
            key += fmt.Sprintf("%d,", sig[band*m.r+row])
        }
        m.bands[key] = append(m.bands[key], id)
    }
}

func (m *MinHashLSH) CandidatePairs() [][2]int {
    seen := map[[2]int]bool{}
    var pairs [][2]int
    for _, ids := range m.bands {
        for i := 0; i < len(ids); i++ {
            for j := i + 1; j < len(ids); j++ {
                p := [2]int{ids[i], ids[j]}
                if p[0] > p[1] {
                    p = [2]int{ids[j], ids[i]}
                }
                if !seen[p] {
                    seen[p] = true
                    pairs = append(pairs, p)
                }
            }
        }
    }
    return pairs
}

func main() {
    m := NewMinHashLSH(10, 4) // k = 40
    m.Add(1, []uint64{1, 2, 3, 4, 5, 6, 7, 8})
    m.Add(2, []uint64{1, 2, 3, 4, 5, 6, 7, 99}) // high Jaccard with id 1
    m.Add(3, []uint64{100, 101, 102, 103})      // disjoint
    fmt.Println(m.CandidatePairs())             // expect [1 2], not [1 3]/[2 3]
}