Classes and OOP in Python

Posted on

Many data structures, like linked lists, stacks, and binary trees, are built from custom objects. Python classes let you define these. You do not need deep object-oriented theory for DSA, just enough to model nodes and structures.

Defining a class#

A class is a blueprint. __init__ runs when you create an instance and sets up its attributes. self refers to the instance being worked on:

 1class Point:
 2    def __init__(self, x, y):
 3        self.x = x
 4        self.y = y
 5
 6    def distance_from_origin(self):
 7        return (self.x ** 2 + self.y ** 2) ** 0.5
 8
 9p = Point(3, 4)
10print(p.x)                        # 3
11print(p.distance_from_origin())   # 5.0

Methods are functions defined inside a class. Their first parameter is always self.

The node: the building block of many structures#

A linked list is a chain of nodes, where each node holds a value and a reference to the next node:

 1class Node:
 2    def __init__(self, value):
 3        self.value = value
 4        self.next = None
 5
 6# Build 1 -> 2 -> 3
 7a = Node(1)
 8b = Node(2)
 9c = Node(3)
10a.next = b
11b.next = c
12
13# Walk the list
14current = a
15while current is not None:
16    print(current.value)   # 1, then 2, then 3
17    current = current.next

That current = current.next traversal pattern appears in almost every linked list algorithm.

A tree node#

A binary tree node is the same idea with two links instead of one:

 1class TreeNode:
 2    def __init__(self, value):
 3        self.value = value
 4        self.left = None
 5        self.right = None
 6
 7root = TreeNode(10)
 8root.left = TreeNode(5)
 9root.right = TreeNode(15)
10print(root.left.value)   # 5

You will use TreeNode throughout the tree lessons in the main course.

A small stack class#

Wrapping behavior in a class keeps your algorithm code clean:

 1class Stack:
 2    def __init__(self):
 3        self.items = []
 4
 5    def push(self, item):
 6        self.items.append(item)
 7
 8    def pop(self):
 9        return self.items.pop()
10
11    def is_empty(self):
12        return len(self.items) == 0
13
14s = Stack()
15s.push(1)
16s.push(2)
17print(s.pop())        # 2
18print(s.is_empty())   # False

A readable representation#

Defining __repr__ makes objects print helpfully, which saves time when debugging:

1class Point:
2    def __init__(self, x, y):
3        self.x, self.y = x, y
4
5    def __repr__(self):
6        return f"Point({self.x}, {self.y})"
7
8print(Point(1, 2))   # Point(1, 2)

You now have the tools to build the structures the curriculum uses. Last stop in the primer: iterators and comprehensions.