Python Skeleton to Use for DSA Interviews
Published in
5 min readJun 3, 2024
Bookmark this one pager that provides you a python skeleton to use for DSA Interviews
Import Useful Modules
- Include the ones that is required
- Google Python Style Guide recommends sorting imports alphabetically and grouping them by their origin (standard library, third-party, and local application/library imports).
- Note : `zip` and `enumerate` are built-in functions so it does not require additional imports.
from bisect import bisect_left, bisect_right
from collections import Counter, defaultdict
from dataclasses import dataclass
from functools import lru_cache
from math import log, ceil, floor
from random import randint
from typing import Any, Dict, List, Tuple
import heapq
Define Classes
- Define the required classes
- Google Python Style Guide recommends
- Class to follow CapWords convention (e.g.MyClassName
).
- Comparison to `None` is performed usingis
operator rather than==
because behaviour of==
may be overridden by__eq__
function.
- Variable and Function name to follow lowercase with words separated by underscores (e.g.helper_function
,my_var
)
- Use `dataclass` decorator if you want to perform vanilla assignment based initialisation. It saves time and keeps it clean.
########################### Quick Custom Classess ##########################
@dataclass
class Event:
"""
Quickly define class with constructor with minimal typing :)
Providing initial value is not required
"""
event_value: int = 0
event_type: str = "start"
event_name: str = "A"
@dataclass(order=True)
class Event:
"""
Make objects comparable using order=True
by automatically adding `__lt__`, `__le__`, `__gt__` and `__ge__`
The sorting will first compare by event_value,
and if event_value values are equal, it will compare by event_type
and ...
"""
event_value: int = 0
event_type: str = "start"
event_name: str = "A"
############################ Heap Implementations ##########################
class MaxHeap:
""" Implementation of Max Heap using heapq """
# Make sure to define appropirate type for data.
def __init__(self, data: List[int] = None):
if data is None:
data = []
self.data = data
heapq.heapify(self.data)
def push(self, n: int):
heapq.heappush(self.data, -n)
def pop(self) -> int:
n = heapq.heappop(self.data)
return -n
def top(self) -> int:
n = self.data[0]
return -n
def __len__(self) -> int:
# Allows to access length in standard way : len(max_heap)
return len(self.data)
class MinHeap:
""" Implementation of Min Heap using heapq """
# Make sure to define appropirate type for data.
def __init__(self, data: List[int] = None):
if data is None:
data = []
self.data = data
heapq.heapify(self.data)
def push(self, n: int):
heapq.heappush(self.data, n)
def pop(self) -> int:
return heapq.heappop(self.data)
def top(self) -> int:
return self.data[0]
def __len__(self) -> int:
# Allows to access length in standard way : len(min_heap)
return len(self.data)
########################## Linked List Implementation ######################
@dataclass
class LinkedListNode:
""" Implementation of Singly Linked List """
val: int
next: "LinkedListNode" = None
@dataclass
class DoublyLinkedListNode:
""" Implementation of Doubly Linked List """
val: int
prev: "DoublyLinkedListNode" = None
next: "DoublyLinkedListNode" = None
########################## Binary Tree Implementation ######################
@dataclass
class TreeNode:
""" Implementation of Node of Binary Tree """
val: int
left: "TreeNode" = None
right: "TreeNode" = None
########################## Trie Implementation ############################
@dataclass
class TrieNode:
""" Implementation of single node of Trie """
children: Dict[str, "TrieNode"] = {}
end_of_word: bool = False
freq: int = 1
class Trie:
""" Implementation of Trie """
def __init__(self):
self.root = TrieNode()
def add(self, word: str):
current_node = self.root
for char in word:
if char not in current_node.children:
current_node.children[char] = TrieNode()
else:
current_node.children[char].freq += 1
current_node = current_node.children[char]
current_node.end_of_word = True
def search_word(self, word: str) -> bool:
current_node = self.root
for char in word:
if char not in current_node.children:
return False
current_node = current_node.children[char]
return current_node.end_of_word
def search_prefix(self, prefix: str) -> bool:
current_node = self.root
for char in word:
if char not in current_node.children:
return False
current_node = current_node.children[char]
else:
return True
##################### DisjointSetUnion Implementation #####################
class DSU:
""" Implementation of DisjointSet Union / UnionFind """
def __init__(self, n: int):
#Initially every node is parent of itself
self.parent = list(range(n))
# Choose rank/size based on appropriate union method
# In most cases, rank should suffice
self.rank = [0]*n # Maintains height of tree
self.size = [1]*n # Maintains size (number of elements) of tree
def find_parent(self, v: int) -> int:
""" Finds a parent of a node along with path compression """
if self.parent[v] != v:
self.parent[v] = self.find(self.parent[v]) # Recursion responsible for path compression
return self.parent[v]
def union(self, u: int, v: int) -> bool:
""" Implementation of union without any optimization """
self.parent[self.find_parent(u)] = self.find_parent(v)
def union_by_rank(self, u: int, v: int) -> bool:
""" Performs union of two disjoint sets that u and v belongs to """
root_u = self.find_parent(u)
root_v = self.find_parent(v)
if root_u == root_v:
# Indicates that u and v is already part of the same set
return False
else:
if self.rank[root_u] > self.rank[root_v]:
self.parent[root_v] = root_u
elif self.rank[root_u] < self.rank[root_v]:
self.parent[root_u] = root_v
else:
self.parent[root_v] = root_u
self.rank[root_u] += 1
return True
def union_by_size(self, u: int, v: int) -> bool:
""" Performs union of two disjoint sets that u and v belongs to """
root_u = self.find_parent(u)
root_v = self.find_parent(v)
if root_u == root_v:
# Indicates that u and v is already part of the same set
return False
else:
if self.size[root_u] > self.size[root_v]:
self.parent[root_v] = root_u
self.size[root_u] += self.size[root_v]
else:
self.parent[root_u] = root_v
self.size[root_v] += self.size[root_u]
return True
##################### Binary Search Implementation #####################
def binary_search(nums: List[int], target: int) -> int:
""" Returns index of target element if found otherwise returns -1 """
lo = 0
hi = len(nums)
while lo < hi:
mid = (lo+hi)//2
if nums[mid] < target:
lo = mid+1
else:
hi = mid
return lo if lo < len(nums) and nums[lo] == target else -1
Define functions
- Define all helper functions that is needed followed by primary function
- I prefer defining helper functions outside primary function rather than using nested functions as this makes helper function individually testable.
def next_number(number: int) -> int:
""" Finds next number by taking summation of square of digits of given number """
sum = 0
while number > 0:
number, digit = divmod(number, 10)
sum += digit**2
return sum
def is_happy_number(number: int) -> bool:
""" Finds whether given number is happy number or not """
slow = number
fast = next_number(number)
while slow != fast and fast != 1:
slow = next_number(number)
fast = next_number(next_number(number))
return number == 1
Define main function
- Responsibilities of main function are
- Prepare inputs required to call primary function
- Call primary function
- Assert result returned by primary function - This is good place to show case that edges cases are covered :)
def main():
assert happy_number(1) == True
assert happy_number(2) == False
assert happy_number(19) == True
Call the main function
- For the sake of completeness (and for the sake of saying that this code can be executed as script), you can call
main()
function
def __name__ == "__main__":
main()
Skeleton
from collections import Counter, defaultdict
from functools import lru_cache
from math import log
from random import randint
from typing import Any, Dict, List, Tuple
import heapq
class HelperClass:
...
def helper_function() -> bool:
...
def helper_function_2() -> bool:
...
def primary_function() -> bool:
...
def main():
assert primary_function() == True
if __name__ == "__main__"
main()