Sorting and Searching Algorithms Questions

Approach: normalize embeddings and compute cosine similarity as dot product between normalized Q and D; for each query return top-k indices. For moderate n (<=1M) and d (<=512) do batched, vectorized matrix multiplications to avoid O(mnd) memory explosion.

import numpy as np

def knn_cosine_bruteforce(Q, D, k, batch_size=1024):
    """
    Q: (m, d) float32
    D: (n, d) float32
    returns: (m, k) indices of top-k neighbors in D per query
    """
    # Normalize in-place (avoid extra memory)
    Q_norm = Q / (np.linalg.norm(Q, axis=1, keepdims=True) + 1e-8)
    D_norm = D / (np.linalg.norm(D, axis=1, keepdims=True) + 1e-8)

    m, d = Q_norm.shape
    n = D_norm.shape[0]
    result = np.empty((m, k), dtype=np.int64)

    for i in range(0, m, batch_size):
        qbatch = Q_norm[i:i+batch_size]                 # (b, d)
        # compute similarity matrix (b, n) via GEMM
        sims = qbatch.dot(D_norm.T)                    # memory: b * n floats
        # get top-k indices per query
        idx_part = np.argpartition(-sims, k-1, axis=1)[:, :k]
        # sort those top-k to get exact order
        rows = np.arange(sims.shape[0])[:, None]
        topk_sorted = idx_part[np.argsort(-sims[rows, idx_part], axis=1)]
        result[i:i+batch_size] = topk_sorted
    return result

Key points:

• Complexity: compute cost O(m * n * d) FLOPs; memory for a single batch similarity matrix is O(batch_size * n).
• Choose batch_size so batch_size * n fits memory (e.g., batch_size = max(1, floor(available_bytes / (4*n)))).
• Use float32 to halve memory vs float64.
• Vectorization: rely on BLAS-backed GEMM (NumPy/Intel MKL, OpenBLAS) or torch.mm on GPU for speed.
• PyTorch variant: move D_norm to GPU once, stream Q batches to GPU, use torch.topk (faster than argpartition+sort).
• Trade-offs:
  • Larger batches increase throughput (better GEMM) but need more RAM/VRAM.
  • Pre-normalizing D costs O(n*d) once; storing normalized D duplicates if you need original vectors.
  • For n ~1M and d=512, storing D_norm ~ 2 GB (1e65124 bytes) — feasible on machines with ~32GB RAM; GPU storage may be limiting.
• Optimizations: use half precision (float16) if acceptable, Product Quantization / IVF for sublinear search as next steps. Edge cases: k>n, zero vectors (handled by eps), empty Q/D.

Requirements & constraints:

• Merge 100,000 sorted runs into one sorted file, memory budget 16 GB, minimize number of passes (disk reads/writes), optimize I/O seeks.

High-level approach:

1. Create as few, as long runs as possible using replacement-selection during initial run generation (memory buffer sized ~M). Replacement-selection typically doubles expected run length versus simple sort, reducing initial run count before final merges.

2. Use a multi-level k-way merge tree: pick a fan-in F per merge level so that number of levels L = ceil(log_F(number_of_runs)). Each level reads from F input streams and writes 1 output stream.

Choosing fan-in (tradeoff):

• Memory per input stream includes a read buffer (B_read), output buffer (B_write), and priority structure overhead. With 16 GB, reserve ~12–14 GB for buffers and tree, keep OS/cache and other processes ~2–4 GB.
• For example, with 100 KB per input buffer you can support ~120k streams in theory (but priority structure overhead matters). Practically choose F in 512–4096 range to keep L small (e.g., F=1024 -> L≈ceil(log_1024(100k))≈2 levels).
• Larger F reduces passes but increases memory and CPU for comparisons and reduces sequential I/O efficiency due to many small reads.

Priority structures:

• Use a tournament tree (loser tree) for k-way selection: it does ~log F comparisons per extracted element but allows efficient replacement; memory O(F) for tree. A binary heap has similar cost (~log F) but requires more comparisons in practice and less cache-friendly when F large. Loser tree is preferred at high fan-in.

I/O optimization:

• Use large sequential read buffers per stream (e.g., 64KB–1MB) to amortize seek cost; align buffers with filesystem block size and prefetch asynchronously.
• Use asynchronous I/O / multiple threads: one thread per group to feed the priority structure, separate writer thread with a sizable output buffer.
• When input runs are on many disks/partitions, stripe reads to exploit parallelism; if on a single HDD, reduce simultaneous seeks by batching runs and using fewer concurrent streams.

Multi-level strategy/example:

• Step A: replacement-selection to cut runs from 100k to ~50k (example).
• Step B: First merge level with F1 = 2048 -> produces ceil(50k/2048) ≈ 25 intermediate runs.
• Step C: Second level merge F2 = 25 -> single final run.
  This yields 2 merge passes after initial generation.

Additional considerations:

• If memory is tight, use more levels with smaller F (e.g., 4–8 levels) but accept more read/write passes.
• Monitor I/O patterns; adapt F per level depending on observed throughput and whether data is on SSDs (favor higher F) or HDDs (favor smaller F and larger buffers).
• Use compression for on-disk runs if CPU cheaper than I/O; reduces I/O at cost of CPU and memory for decompression buffers.

Result: By combining replacement-selection to reduce initial runs, using loser trees for high fan-in merges, and choosing F to keep levels to 2–3 while fitting buffer+tree memory into 16 GB, you minimize passes and optimize end-to-end throughput.

std::nth_element semantics and complexity

• Semantics: std::nth_element(first, nth, last) reorders the range in-place so that the element at position nth is the same as it would be in a fully sorted sequence. All elements before nth are <= that element and all elements after are >= it (no guarantee of order within those partitions). It's a partial sort, not stable.
• Complexity: average O(n) time and O(1) extra space (in-place), implemented via Quickselect/partitioning. Worst-case is O(n^2) without safeguards, but standard-library implementations often use introspective techniques to avoid pathological worst-cases.

Equivalent behavior in Python

• numpy.partition implements the same idea (Quickselect in C): average O(n) and operates on ndarray, returning a new array or doing in-place partitioning.
• For streaming or memory-limited scenarios, heap-based methods (heapq.nsmallest / heapq.nlargest) run in O(n log k) time and O(k) extra space — preferable when k << n or when input is an iterator.

Python sketches

Using numpy.partition (fast for large arrays on CPU/GPU-capable arrays):

import numpy as np

arr = np.random.rand(1_000_000)
k = 10
# returns array where first k elements are the k smallest in arbitrary order
part = np.partition(arr, k)[:k]
# to get them sorted:
topk_sorted = np.sort(part)  # O(k log k)

Using heapq for streaming / small k:

import heapq
def topk_iterator(iterable, k):
    return heapq.nlargest(k, iterable)  # O(n log k), uses O(k) mem

When to prefer nth_element-style partial ordering in ML

• Top-k inference: you only need the top-k logits/indices (argmax-k) for sampling, beam search or returning predictions — full sort is wasteful.
• Beam search / decoding: keep top-k candidates per step; partial ordering suffices and reduces compute.
• Non-maximum suppression / object proposal pruning: select top-k proposals by score before expensive IoU computations.
• Candidate generation/pruning in retrieval or approximate nearest neighbor: reduce the candidate pool quickly.
• Feature selection / quantile computations: median-of-array or k‑th quantiles without full sorting.

Trade-offs and practical notes

• Use numpy.partition when working with large dense tensors (vectorized, C-optimized).
• Use heapq for streaming inputs or when k is tiny relative to n.
• On GPUs, libraries (PyTorch: torch.topk uses partial selection) provide efficient top-k primitives — prefer those for GPU tensors to avoid host-device transfers.
• If you need stable order or complete ordering inside the top-k, sort the k results (O(k log k)) after partitioning.

Approach:
Use LSD radix sort with base 256 (one byte per pass). For 32-bit non-negative ints we need at most 4 stable passes. Each pass is a stable counting sort on the byte at position p (0 = least significant). To keep stability, build cumulative counts then place elements into an output array scanning input left→right.

Code implementation (pure Python, works with Python ints but for performance and memory use array('I') or numpy.uint32 as recommended below):

from array import array

def lsd_radix_sort_256(arr):
    """
    LSD radix sort base 256 for non-negative 32-bit ints.
    Uses at most 4 passes. Stable counting sort per pass.
    Works in O(4n) time = O(n).
    Input: array('I') or list of non-negative ints.
    Returns sorted list/array of same type.
    """
    if not arr:
        return arr

    n = len(arr)
    # Use arrays for lower memory; if input is list, convert to array for efficiency
    is_array = isinstance(arr, array)
    if not is_array:
        a = array('I', arr)
    else:
        a = arr

    out = array('I', [0]) * n
    MASK = 0xFF

    for pass_no in range(4):  # 4 bytes for 32-bit
        counts = [0] * 256
        shift = pass_no * 8

        # Count occurrences of each byte value
        for x in a:
            counts[(x >> shift) & MASK] += 1

        # Cumulative counts -> starting index for each bucket
        total = 0
        for i in range(256):
            c = counts[i]
            counts[i] = total
            total += c

        # Place elements into out preserving order (stable): left->right
        for x in a:
            b = (x >> shift) & MASK
            out[counts[b]] = x
            counts[b] += 1

        # swap buffers for next pass
        a, out = out, a

    # After 4 passes 'a' holds sorted data
    if is_array:
        # if original was array, copy back
        for i in range(n):
            arr[i] = a[i]
        return arr
    else:
        return list(a)

Key concepts and reasoning:

• LSD processes least significant byte first; stability per pass ensures global order by lower bytes is preserved when sorting by higher bytes.
• Counting sort is ideal for small fixed radix (256): counts array size is constant (256), so per-pass work is O(n).
• We swap input/output buffers each pass to avoid repeated allocations per element.

Time and space complexity:

• Time: O(4n) = O(n) for 32-bit ints (4 passes).
• Space: O(n) extra for output buffer + O(256) for counts. If using Python lists of PyObjects, memory per element is large; using array('I') or numpy.uint32 reduces memory to ~4 bytes/element.

Edge cases:

• Empty input, already-sorted input, many duplicates — all handled.
• If integers can exceed 32 bits, increase passes accordingly.
• Negative numbers: current code assumes non-negative; to support negatives, map to unsigned via offset or handle sign byte specially.

Memory/time trade-offs and optimizations for large arrays (n up to 50M):

• Python-level lists of ints are memory-inefficient (28+ bytes per int) — infeasible for 50M. Use compact representations:
  • array('I') or numpy.uint32: ~4 bytes per element → 200MB for 50M elements.
  • memory-map arrays (numpy.memmap) to avoid loading entire data into RAM.
• Minimize Python overhead:
  • Implement core loop in C (use numpy vectorized ops, Cython, or Numba) to eliminate Python interpreter cost.
  • Process in chunks with memmap or streaming if RAM is limited; maintain global counts per chunk then redistribute (more complex to maintain stability).
• Parallelization:
  • Parallel counting: split input into chunks, compute local counts, reduce to global counts, compute offsets, then write per-chunk into output (requires coordination of offsets). Best implemented in C/C++ or using multiprocessing with shared memory for large datasets.
• Allocation strategy:
  • Reuse two preallocated buffers (as in code) to avoid repeated allocations.
  • For very large n, prefer memory-mapped output or in-place strategies in lower-level languages.
• Practical recommendation for 50M:
  • Use numpy.uint32 with memmap and implement counting placement in C/Cython/Numba or use external sort library. This gives manageable memory footprint and acceptable runtime.

Alternative improvements:

• Use radix = 65536 (2 bytes) to reduce passes to 2, but counts size becomes 65536 (256KB), still small — trade-off: larger counts array fits CPU cache? Two passes may reduce memory traffic but counts array size rises; choose based on cache behavior and memory.
• For unknown max value, compute max to limit number of passes.

This implementation is stable, predictable O(n) time, and practical if you use compact numeric arrays and move heavy loops out of pure Python for very large n.

Non-comparison sorts exploit key structure instead of pairwise comparisons. Briefly:

Counting sort

• Keys: non-negative integers (or mapped to small integer range) with range U.
• Stable: yes (if implemented by scanning cumulative counts).
• Time: O(n + U). Space: O(n + U).
• Use cases: histogramming integer features, converting categorical labels to counts, fast preprocessing for one-hot / embedding frequencies.
• When better: when U = O(n) (or small constant factor) so O(n) beats O(n log n).

Radix sort

• Keys: integers or fixed-length digit strings; can handle floats by bitwise reinterpretation or separate sign/exponent/mantissa processing.
• Stability: stable if each digit pass uses a stable subroutine (usually counting sort).
• Time: O(d (n + b)) where d = digits, b = base per digit. For w-bit integers, O((w / log b) * (n + b)). Space: O(n + b).
• Use cases: sorting 32/64-bit integer IDs, large index arrays, batching by hashed keys, lexicographic sorting of fixed-length tokens.
• When better: when key width is bounded and d is small → linear-time practical advantage for large n.

Bucket sort

• Keys: typically real numbers in a known range (e.g., [0,1)); works for floats if distribution is roughly uniform or buckets sized appropriately.
• Stability: can be stable depending on bucket implementation and internal sort.
• Time: average O(n + k) where k = number of buckets (with O(n²) worst-case if buckets uneven and internal sort is quadratic). Space: O(n + k).
• Use cases: approximate quantile computation, stratified batching, distributing items for parallel processing, smoothing continuous features.
• When better: when keys are uniformly distributed so buckets remain small — gives near-linear performance.

Practical notes for AI/data pipelines

• Use counting/radix for massive integer keys (IDs, binned features, discrete vocab indices) to build frequency tables, sharding maps, or sorted indices quickly.
• Bucket sort is useful for bucketing continuous features for stratified sampling, minibatch construction, or approximate quantiles.
• They outperform comparison sorts when keys are low-entropy, bounded, or fixed-size (leading to O(n) behavior) and when memory for auxiliary arrays is acceptable.