11. Classes & Dataclasses

Tutorial

DSA primarily uses classes for node definitions. Know __lt__ for heapq compatibility and @dataclass for clean boilerplate-free classes.

# ══ DSA NODE CLASSES ═══════════════════════════════════════════
class ListNode:
    def __init__(self, val: int = 0, next: "ListNode | None" = None):
        self.val  = val
        self.next = next
    def __repr__(self) -> str:
        vals, cur = [], self
        while cur:
            vals.append(str(cur.val)); cur = cur.next
        return " → ".join(vals)

class TreeNode:
    def __init__(self, val: int = 0,
                 left:  "TreeNode | None" = None,
                 right: "TreeNode | None" = None):
        self.val   = val
        self.left  = left
        self.right = right

# Graph node (sometimes used)
class Node:
    def __init__(self, val: int, neighbors: list["Node"] | None = None):
        self.val       = val
        self.neighbors = neighbors if neighbors is not None else []

# ══ __lt__ FOR HEAPQ COMPATIBILITY ════════════════════════════
# heapq requires elements to be comparable
class State:
    def __init__(self, cost: int, node: int):
        self.cost = cost
        self.node = node
    def __lt__(self, other: "State") -> bool:
        return self.cost < other.cost   # heapq uses < operator

import heapq
heap = [State(5, "B"), State(1, "A")]
heapq.heapify(heap)
top = heapq.heappop(heap)   # State(1, "A")

# ══ DATACLASS ══════════════════════════════════════════════════
from dataclasses import dataclass, field

@dataclass
class Point:
    x: int
    y: int
    # auto-generates __init__, __repr__, __eq__

@dataclass(order=True)  # also generates __lt__, __le__, __gt__, __ge__
class Edge:
    weight: int          # compared first (field order matters)
    src: int = field(compare=False)  # excluded from comparison
    dst: int = field(compare=False)

# Push Edge directly into heap (order=True adds __lt__)
edges = [Edge(3, 0, 1), Edge(1, 0, 2)]
heapq.heapify(edges)

# @dataclass(frozen=True) — immutable, hashable (can be dict key / set element)
@dataclass(frozen=True)
class Coord:
    row: int
    col: int

visited: set[Coord] = set()
visited.add(Coord(0, 0))   # works because frozen=True makes it hashable

# ══ __slots__ (memory optimisation for many small objects) ════
class CompactNode:
    __slots__ = ("val", "next")   # prevents __dict__ creation
    def __init__(self, val: int):
        self.val  = val
        self.next = None

# ══ INHERITANCE (occasionally needed) ══════════════════════════
class Animal:
    def __init__(self, name: str):
        self.name = name
    def speak(self) -> str:
        return f"{self.name} makes a sound."

class Dog(Animal):
    def speak(self) -> str:             # override
        return f"{self.name} barks!"
    def fetch(self) -> str:
        return f"{self.name} fetches!"

d = Dog("Rex")
d.speak()    # "Rex barks!"
d.fetch()    # "Rex fetches!"
isinstance(d, Dog)     # True
isinstance(d, Animal)  # True — Dog is a subclass of Animal

Reference

Dunder	Purpose	DSA use
__init__	Constructor	Node initialisation
__repr__	Debug string	Visualise linked list/tree
__lt__	Less-than comparison	heapq compatibility
__eq__	Equality	Counter, set membership
__hash__	Hashable	Set/dict key (auto with frozen=True)
@dataclass	Auto-generate __init__, __repr__, __eq__	Cleaner node/edge classes
order=True	Also generate __lt__ etc.	heapq without manual __lt__
frozen=True	Immutable + hashable	Dict key, set element
field(compare=False)	Exclude from ordering	Priority + data in heap
__slots__	No __dict__, less memory	Large graphs
class Dog(Animal)	Inheritance	Override methods; isinstance() checks subclass

---