Skip to content

Data Structures Essentials for ML

Overview

Python's built-in data structures are not the fastest, but they're the foundation for: - NumPy arrays (built on lists conceptually) - PyTorch tensors (batched operations on lists) - DataLoaders (list-based batching) - Model configs (nested dicts)

Understanding their implementation helps optimize ML pipelines.

Lists: Sequential Storage

Core Characteristics

# Lists are dynamic arrays
data = [1, 2, 3]
print(type(data))  # <class 'list'>

# O(1) append (amortized), O(1) indexing
data.append(4)     # Appends to internal C array
data[0]            # Direct memory access

# But O(n) insert at beginning
data.insert(0, 0)  # Shifts all elements

How Lists Work Internally

Python List = [Reference, Reference, Reference, ...]
              ↓          ↓          ↓
            Object    Object    Object
              (int)    (int)     (int)
  • Dynamic sizing: Allocates extra capacity to avoid reallocations
  • Heterogeneous: Can store any type (overhead vs NumPy's homogeneous)
  • Reference semantics: Stores pointers, not values

ML Pattern: DataLoader with Lists

from torch.utils.data import DataLoader, TensorDataset
import torch

# Lists are converted to tensors
images = [torch.randn(3, 224, 224) for _ in range(1000)]
labels = [i % 10 for i in range(1000)]

# DataLoader batches lists → tensors
dataset = TensorDataset(torch.stack(images), torch.tensor(labels))
loader = DataLoader(dataset, batch_size=32, shuffle=True)

for batch_images, batch_labels in loader:
    print(batch_images.shape)  # (32, 3, 224, 224)

Dictionaries: Hash Tables

Core Characteristics

# Dicts are hash tables
config = {"learning_rate": 0.001, "batch_size": 32}

# O(1) average lookup, insertion, deletion
config["learning_rate"]  # O(1)
config["new_key"] = 0.1  # O(1)
del config["batch_size"]  # O(1)

# Python 3.7+: Insertion order preserved
# Before 3.7: Order was arbitrary

How Dicts Work Internally

Hash Table:
[None] → (key, value)
[Hash(key) % size] → (key, value)
...

Collision handling: Open addressing (probe for next empty slot)

ML Framework Pattern: Model Configs

import torch
from dataclasses import dataclass
from typing import Dict, Any

# HuggingFace-style config (dict underneath)
class ModelConfig(dict):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.__dict__ = self  # Enable attribute access

config = ModelConfig(
    hidden_size=768,
    num_attention_heads=12,
    intermediate_size=3072,
    num_hidden_layers=12
)

# Both work
config['hidden_size']      # Dict access
config.hidden_size         # Attribute access

# Easy serialization
import json
json_str = json.dumps(config)

State Dicts in PyTorch

model = torch.nn.Linear(10, 5)

# state_dict is OrderedDict
state = model.state_dict()
print(type(state))  # odict

# Fast lookup of parameters
for name, param in state.items():
    print(f"{name}: {param.shape}")
    # weight: torch.Size([5, 10])
    # bias: torch.Size([5])

# Efficient serialization
torch.save(state, "model.pth")  # Uses pickle on OrderedDict

Sets: Hash-Based Collections

Core Characteristics

# Sets are hash tables with no values (only keys)
unique_ids = {1, 2, 3, 1, 2}  # {1, 2, 3}

# O(1) membership testing
if 1 in unique_ids:           # Fast!
    print("Found")

# Set operations
ids_a = {1, 2, 3}
ids_b = {2, 3, 4}

ids_a & ids_b  # Intersection: {2, 3}
ids_a | ids_b  # Union: {1, 2, 3, 4}
ids_a - ids_b  # Difference: {1}

ML Pattern: Deduplication

# Remove duplicate samples in dataset
def deduplicate_dataset(samples):
    """Remove duplicate samples efficiently."""
    seen = set()
    unique = []

    for sample in samples:
        sample_id = sample['id']  # O(1) lookup
        if sample_id not in seen:
            unique.append(sample)
            seen.add(sample_id)

    return unique

# For large-scale dedup
large_ids = set(torch.load("all_ids.pt").tolist())  # Fast lookup

Tuples: Immutable Sequences

Core Characteristics

# Tuples are immutable lists
shape = (3, 224, 224)

# Immutability benefits
# 1. Hashable (can be dict keys)
config = {(3, 224, 224): "image_model"}

# 2. Cannot be accidentally modified
def process(data_shape: tuple) -> None:
    # data_shape[0] = 4  # TypeError: tuples don't support item assignment

# 3. Slightly faster than lists (no dynamic sizing)

Tuple Unpacking

# Used everywhere in PyTorch
batch_size, channels, height, width = (32, 3, 224, 224)

# Function returns tuple
def get_tensor_info(tensor):
    return tensor.shape, tensor.dtype, tensor.device

shape, dtype, device = get_tensor_info(torch.zeros(3, 4))

# Named tuples (more readable)
from typing import NamedTuple

class TensorInfo(NamedTuple):
    shape: tuple
    dtype: torch.dtype
    device: torch.device

info = TensorInfo(shape=(3, 4), dtype=torch.float32, device=torch.device('cpu'))
print(info.shape)  # (3, 4)

Pattern: Return Multiple Values

def forward_with_cache(
    input: torch.Tensor,
    state: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return output and new state."""
    output = input + state
    new_state = state * 0.9
    return output, new_state

# Easy unpacking
out, new_state = forward_with_cache(x, state)

Nested Data Structures

Lists of Dicts (Most Common in ML)

# Dataset format (HuggingFace, PyTorch)
dataset = [
    {"text": "Hello world", "label": 0, "id": 1},
    {"text": "Good morning", "label": 1, "id": 2},
    {"text": "Hi there", "label": 0, "id": 3},
]

# Access patterns
first_sample = dataset[0]
first_text = dataset[0]["text"]  # O(1) indexing, then O(1) dict lookup

Dicts of Lists (Batch Processing)

# PyTorch DataLoader batches into this format
batch = {
    "input_ids": torch.tensor([[101, 2054, 2088], ...]),  # (batch_size, seq_len)
    "attention_mask": torch.tensor([[1, 1, 1], ...]),
    "labels": torch.tensor([0, 1, 0, ...]),  # (batch_size,)
}

# Efficient batch operations
batch_size = batch["input_ids"].shape[0]  # O(1)
seq_len = batch["input_ids"].shape[1]     # O(1)

Performance Characteristics

Time Complexity

Operation List Dict Set Tuple
Lookup O(n) O(1)* O(1)* O(n)
Insert O(n) O(1)* O(1)* N/A
Delete O(n) O(1)* O(1)* N/A
Iteration O(n) O(n) O(n) O(n)

*Average case, assuming good hash function

Memory Characteristics

import sys

# Lists have overhead
lst = [1, 2, 3]
print(sys.getsizeof(lst))  # ~88 bytes for 3-element list
print(sys.getsizeof(lst) + sys.getsizeof(1) * 3)  # Add object sizes

# Dicts have more overhead (for hash table)
d = {"a": 1, "b": 2}
print(sys.getsizeof(d))  # ~240 bytes for 2-element dict

# Sets similar to dicts (no values)
s = {1, 2}
print(sys.getsizeof(s))  # ~216 bytes

# Tuples more efficient
t = (1, 2, 3)
print(sys.getsizeof(t))  # ~48 bytes for 3-element tuple

For ML: NumPy/PyTorch arrays are much more efficient than Python collections.

Copy Semantics

Critical for avoiding bugs in ML:

Shallow vs Deep Copy

import copy

# Lists: shallow copy by default
original = [[1, 2], [3, 4]]
shallow = original.copy()
deep = copy.deepcopy(original)

# Modify nested list
original[0][0] = 999

print(shallow[0][0])  # 999 (shallow copy affected!)
print(deep[0][0])     # 1 (deep copy unaffected)

# Dicts: same behavior
config = {"model": {"hidden_size": 768}}
config_copy = config.copy()
config["model"]["hidden_size"] = 512
print(config_copy["model"]["hidden_size"])  # 512 (shallow!)

ML Pattern: Config Copying

import copy
from dataclasses import dataclass

@dataclass
class ModelConfig:
    hidden_size: int = 768
    num_layers: int = 12

def create_model_variants(base_config):
    """Create model configs with different sizes."""
    configs = []

    for size_factor in [1, 2, 4]:
        # WRONG: config = base_config  # All share same object!
        # CORRECT: Use copy
        config = copy.deepcopy(base_config)
        config.hidden_size *= size_factor
        configs.append(config)

    return configs

Comprehensions: Functional List Construction

# List comprehension (faster than append loops)
squares = [x**2 for x in range(10)]

# With condition
even_squares = [x**2 for x in range(10) if x % 2 == 0]

# Nested comprehension
matrix = [[i*10 + j for j in range(3)] for i in range(3)]

# Dict comprehension
config = {f"layer_{i}": 768 for i in range(12)}

# Set comprehension
unique = {x % 3 for x in range(100)}

# Generator expression (lazy, memory efficient)
sum_of_squares = sum(x**2 for x in range(1000000))  # No intermediate list

Practical ML Patterns

Pattern 1: Batch Preparation

# Convert list of samples to batch dict
samples = [
    {"text": "hello", "label": 0},
    {"text": "world", "label": 1},
]

batch = {
    "texts": [s["text"] for s in samples],
    "labels": [s["label"] for s in samples],
}

Pattern 2: Model Architecture as Nested Dict

architecture = {
    "encoder": {
        "type": "transformer",
        "hidden_size": 768,
        "num_layers": 12,
        "attention": {
            "num_heads": 12,
            "dropout": 0.1,
        }
    },
    "decoder": {
        "type": "linear",
        "output_size": 10,
    }
}

Pattern 3: Efficient Tensor Batching

# SLOW: Append to list, then convert
tensors = []
for data in dataset:
    tensors.append(process(data))
batch_tensor = torch.stack(tensors)  # Allocates new memory

# FAST: Pre-allocate, fill in place
batch_tensor = torch.zeros(len(dataset), 3, 224, 224)
for i, data in enumerate(dataset):
    batch_tensor[i] = process(data)  # In-place modification

Trade-offs: Python vs NumPy/PyTorch

Aspect Python Collections NumPy/PyTorch
Speed Slow (Python bytecode) Fast (vectorized C)
Memory High overhead Efficient
Flexibility Very flexible Type-specific
Use Case Configs, batching Computation

Best Practice: Use Python collections for structure, NumPy/PyTorch for computation.

Summary

  • Lists = Dynamic arrays, used for batching
  • Dicts = Fast lookup, used for configs and state_dicts
  • Sets = Deduplication, membership testing
  • Tuples = Immutable, used for shapes and return values
  • Comprehensions = Efficient list/dict construction
  • Shallow copy = Gotcha when copying nested structures