Python Coding and Data Structures Questions

import hashlib, math, heapq
from collections import defaultdict, Counter, deque

class CountMinSketch:
    def __init__(self, eps=1e-3, delta=1e-3):
        self.w = int(math.ceil(math.e / eps))
        self.d = int(math.ceil(math.log(1/delta)))
        self.count = [ [0]*self.w for _ in range(self.d) ]
        self.seed = [i*131 for i in range(self.d)]
    def _hash(self, x, i):
        h = int(hashlib.md5((str(self.seed[i])+x).encode()).hexdigest(),16)
        return h % self.w
    def add(self, x, c=1):
        for i in range(self.d):
            self.count[i][self._hash(x,i)] += c
    def estimate(self,x):
        return min(self.count[i][self._hash(x,i)] for i in range(self.d))
    def merge(self, other):
        assert self.d==other.d and self.w==other.w
        for i in range(self.d):
            for j in range(self.w):
                self.count[i][j] += other.count[i][j]

class SpaceSaving:
    # keeps top-M candidates with approximate counts
    def __init__(self, M):
        self.M = M
        self.c = {}  # key -> count
    def add(self, key, incr=1):
        if key in self.c:
            self.c[key] += incr
        elif len(self.c) < self.M:
            self.c[key] = incr
        else:
            # replace min element
            kmin = min(self.c, key=lambda k: self.c[k])
            self.c[key] = self.c.pop(kmin) + incr
    def merge(self, other):
        for k,v in other.c.items():
            self.add(k,v)

def ngram_stream(fileobj, n=2):
    window = deque()
    for line in fileobj:
        tokens = line.strip().split()
        for t in tokens:
            window.append(t)
            if len(window) > n:
                window.popleft()
            if len(window) == n:
                yield ' '.join(window)
        # reset across lines if desired:
        window.clear()

def process_shard(path, n=2, cms_params=(1e-3,1e-3), topM=1000):
    cms = CountMinSketch(*cms_params)
    ss = SpaceSaving(topM)
    with open(path, 'r', encoding='utf8') as f:
        for ng in ngram_stream(f, n=n):
            cms.add(ng)
            ss.add(ng)
    return cms, ss

def merge_results(list_of_pairs):
    cms0, ss0 = list_of_pairs[0][0], list_of_pairs[0][1]
    for cms, ss in list_of_pairs[1:]:
        cms0.merge(cms)
        ss0.merge(ss)
    return cms0, ss0

def top_k_from_merged(cms, ss, k):
    # candidate keys from SpaceSaving
    heap = []
    for key in ss.c.keys():
        est = cms.estimate(key)
        if len(heap) < k:
            heapq.heappush(heap, (est, key))
        else:
            if est > heap[0][0]:
                heapq.heapreplace(heap, (est, key))
    return sorted(heap, reverse=True)

# Example usage: process shards in parallel (e.g., multiprocessing) returning list of (cms,ss) pairs,
# then call merge_results and top_k_from_merged on merged objects.

import numpy as np
from multiprocessing import Pool
from functools import partial

def _chunk_pairs(n, chunk):
    ranges = [(i, min(n, i+chunk)) for i in range(0, n, chunk)]
    for a_idx, (ai, aj) in enumerate(ranges):
        for b_idx, (bi, bj) in enumerate(ranges):
            if b_idx < a_idx: 
                continue  # only upper triangle including diagonal
            yield (ai, aj, bi, bj)

def pairwise_euclidean(X, chunk=1024, workers=0, dtype=np.float32):
    """
    X: (N, D) array (fits in memory)
    chunk: rows per chunk
    workers: 0 => no multiprocessing, >0 => number of worker processes
    dtype: np.float32 or np.float64
    returns: (N, N) distance matrix (dtype)
    """
    X = np.asarray(X, dtype=dtype)
    N = X.shape[0]
    # precompute squared norms
    norms = np.sum(X * X, axis=1, dtype=dtype)  # shape (N,)
    D = np.empty((N, N), dtype=dtype)
    D.fill(0)

    def process_pair(pair):
        ai, aj, bi, bj = pair
        A = X[ai:aj]           # (a, D)
        B = X[bi:bj]           # (b, D)
        # compute cross term using BLAS via dot
        K = A.dot(B.T)        # (a, b)
        # broadcasting norms
        sub = norms[ai:aj, None] + norms[bi:bj][None, :] - 2.0 * K
        # numerical safety: clip tiny negatives to zero
        np.clip(sub, 0, None, out=sub)
        # sqrt to get distances
        np.sqrt(sub, out=sub)
        return (ai, aj, bi, bj, sub)

    pairs = list(_chunk_pairs(N, chunk))
    if workers and len(pairs)>0:
        with Pool(workers) as p:
            for ai, aj, bi, bj, sub in p.imap_unordered(process_pair, pairs):
                D[ai:aj, bi:bj] = sub
                if ai!=bi:
                    D[bi:bj, ai:aj] = sub.T
    else:
        for pair in pairs:
            ai, aj, bi, bj, sub = process_pair(pair)
            D[ai:aj, bi:bj] = sub
            if ai!=bi:
                D[bi:bj, ai:aj] = sub.T
    return D

import numpy as np

def pairwise_cosine_similarity(A, B, batch_size=1024, eps=1e-12, dtype=None):
    """
    Compute pairwise cosine similarity between A (m x d) and B (n x d) in a memory-efficient way.
    Returns an (m x n) array if m*n fits memory; otherwise yields rows incrementally (generator).
    By default returns full matrix; to avoid big output, iterate using yield_rows=True pattern.
    """
    A = np.asarray(A)
    B = np.asarray(B)
    if dtype is not None:
        A = A.astype(dtype, copy=False)
        B = B.astype(dtype, copy=False)

    m, d = A.shape
    n, d2 = B.shape
    if d != d2:
        raise ValueError("Dimension mismatch between A and B")

    # Precompute B norms in chosen dtype for stability
    compute_dtype = np.float64 if A.dtype == np.float32 else A.dtype
    B_ld = B.astype(compute_dtype, copy=False)
    B_norms = np.linalg.norm(B_ld, axis=1)
    B_norms = np.maximum(B_norms, eps)  # avoid division by zero
    B_normalized = (B_ld.T / B_norms).T  # n x d

    # Allocate output (be mindful of size; caller can set batch_size to avoid huge allocation)
    result = np.empty((m, n), dtype=compute_dtype)

    for i in range(0, m, batch_size):
        Ai = A[i : i + batch_size].astype(compute_dtype, copy=False)   # bs x d
        A_norms = np.linalg.norm(Ai, axis=1)
        A_norms = np.maximum(A_norms, eps)
        Ai_normalized = (Ai.T / A_norms).T  # bs x d

        # bs x d  dot  d x n  -> bs x n
        sims = Ai_normalized.dot(B_normalized.T)
        # numerical rounding can push values slightly beyond [-1,1]
        np.clip(sims, -1.0, 1.0, out=sims)
        result[i : i + batch_size] = sims

    return result

import json, struct, numpy as np, tempfile, os
MAGIC=b'NPYBUNDLE'
VERSION=1

def write_bundle(path, arrays):  # arrays: dict[name]=(ndarray, metadata, compressor=None)
    with open(path+'.tmp','wb') as f:
        f.write(MAGIC)
        f.write(struct.pack('B', VERSION))
        f.write(struct.pack('<Q', 0))  # placeholder for header offset
        entries=[]
        for name,(arr,meta,comp) in arrays.items():
            offset=f.tell()
            data=arr.tobytes(order='C')
            if comp=='zstd':
                import zstandard as zstd
                data=zstd.ZstdCompressor().compress(data)
            length=len(data)
            f.write(data)
            entries.append({'name':name,'dtype':arr.dtype.str,'shape':arr.shape,'order':'C',
                            'offset':offset,'length':length,'compressed': bool(comp),'compressor':comp,'meta':meta})
        header_bytes=json.dumps({'entries':entries}).encode()
        header_offset=f.tell()
        f.write(header_bytes)
        f.flush()
    # write header_offset into placeholder
    with open(path+'.tmp','r+b') as f:
        f.seek(len(MAGIC)+1)
        f.write(struct.pack('<Q', header_offset))
    os.replace(path+'.tmp', path)

def open_array(path, name):
    import struct, json, numpy as np, tempfile
    with open(path,'rb') as f:
        magic=f.read(len(MAGIC)); assert magic==MAGIC
        version=struct.unpack('B',f.read(1))[0]
        header_off=struct.unpack('<Q',f.read(8))[0]
        f.seek(header_off); header=json.load(f)
    entry=next(e for e in header['entries'] if e['name']==name)
    if not entry['compressed']:
        return np.memmap(path, dtype=np.dtype(entry['dtype']), mode='r',
                         offset=entry['offset'], shape=tuple(entry['shape']), order=entry['order'])
    else:
        # decompress to temp file and memmap
        with open(path,'rb') as f:
            f.seek(entry['offset']); data=f.read(entry['length'])
        if entry['compressor']=='zstd':
            import zstandard as zstd
            data=zstd.ZstdDecompressor().decompress(data)
        tf=tempfile.NamedTemporaryFile(delete=False)
        tf.write(data); tf.flush(); tf.close()
        return np.memmap(tf.name, dtype=np.dtype(entry['dtype']), mode='r',
                         offset=0, shape=tuple(entry['shape']), order=entry['order'])

class Node:
    def __init__(self, k, v):
        self.k = k; self.v = v
        self.prev = None; self.next = None

class LRUCache:
    def __init__(self, capacity):
        self.capacity = capacity
        self.map = {}  # key -> Node
        # dummy head/tail to simplify edge ops
        self.head = Node(None, None)
        self.tail = Node(None, None)
        self.head.next = self.tail
        self.tail.prev = self.head

    # helper: remove node from list
    def _remove(self, node):
        p, n = node.prev, node.next
        p.next = n; n.prev = p

    # helper: insert node right after head (most recent)
    def _add_to_head(self, node):
        node.next = self.head.next
        node.prev = self.head
        self.head.next.prev = node
        self.head.next = node

    def get(self, key):
        node = self.map.get(key)
        if not node:
            return -1
        # mark recently used
        self._remove(node)
        self._add_to_head(node)
        return node.v

    def put(self, key, value):
        node = self.map.get(key)
        if node:
            node.v = value
            self._remove(node)
            self._add_to_head(node)
            return
        if len(self.map) >= self.capacity:
            # evict least recently used (tail.prev)
            lru = self.tail.prev
            self._remove(lru)
            del self.map[lru.k]
        new = Node(key, value)
        self.map[key] = new
        self._add_to_head(new)