Python Skeleton to Use for DSA Interviews

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 using is 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()