Python Algorithm Notes

STL

list

l = [5, 4, 3, 2, 1]
l2 = [i for i in range(5, 0, -1)] # listcomps
pair = [a, 1]
l.insert(2,33)
l.pop()
# here 5 represents the first element
l.index(5)
l.count(5)
l.remove(5)
sorted(l, key=lambda x: (-x[0], x[1]))  # descend at 1st dim, ascend at 2nd

Be aware of order in listcomps:

[leaf for branch in tree for leaf in branch]
It unrolls like:
for branch in tree:
    for leaf in branch:
        yield leaf

dict/set

d = dict()
d2 = {k:ord(k) for k in 'dict'}

st1, st2 = set([1, 2]), set('12') # st2 is {'1', '2'}
st3 = {k for k in 'abracadabra' if x not in 'abc'}
set_or = st1|st2
set_and = a&b
set_diff = a-b
set_nor = a^b # == (a-b)|(b-a)
st3 = {}

Unfortunately, Python’s built-in library does not provide an ordered data structure that can be added, deleted, modified, and checked in $O(logn)$ time.

collections

from collections import deque
d = deque('ghi') 
d.append('j')
d.appendleft('f')
d.pop()
d.popleft()
left, right = d[0], d[-1]

from collections import Counter
Counter('abracadabra').most_common(3) # [('a', 5), ('b', 2), ('r', 2)]

from collections import defaultdict
dd = defaultdict(list)
dd['a'].append(2)

string

s = 'skjadhfsiaf9823r'
s.count('sa', beg=0, end=len(s))
s.find('sa', beg=0, end=len(string)) # return the first index or -1
s.rfind('sa', beg=0, end=len(string))
s.index('sa', beg=0, end=len(string))# return the first index or throw an exception
'ha{}ha{}'.format(1, 2) #ha1ha2

# return False if use '' for following judgement
s.isalpha()
s.isdigit()
s.islower()
s.isupper() 

s = s.lower()
s = s.upper()
s = s.replace('sk', 'ks', num=s.count(str1))
l = s.split('sk', num=s.count(str1))

heap/priority_queue

import heapq
nums = [2, 3, 1, 7, 4]
heapq.heapify(nums)
heapq.heappush(nums, 8)
h = heapq.heappop(nums)

regex

improt re
text = "He was carefully disguised but captured quickly by police."
re.findall(r"\w+ly\b", text)  # ['carefully', 'quickly']
for m in re.finditer(r"\w+ly\b", text):
    print('%d-%d: %s' % (m.start(), m.end(), m.group(0))) # 7-16: carefully 40-47: quickly

Tips

1. search direction in grid:

for nx, ny in ((x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)):

2. Iterate letters:

list(map(chr, range(ord('a'), ord('z') + 1)))

3. Fast power with modulo: pow(a, b, MOD).

High Frequency Templates

Union-Find Set

class UnionFind:
    def __init__(self, n: int) -> None:
        self.root = [i for i in range(n)]
        self.size = [1] * n
        self.part = n
        return
'''
    def find(self, x):
        lst = []
        while x != self.root[x]:
            lst.append(x)
            x = self.root[x]
        for w in lst:
            self.root[w] = x
        return x
'''
    def find(self, x):
        if x != self.root[x]:
        	self.root[x] = self.find(self.root[x])
        return self.root[x]

    def union(self, x, y):
        root_x = self.find(x)
        root_y = self.find(y)
        if root_x == root_y:
            return False
        if self.size[root_x] >= self.size[root_y]:
            root_x, root_y = root_y, root_x
        self.root[root_x] = root_y
        self.size[root_y] += self.size[root_x]
        self.size[root_x] = 0
        self.part -= 1
        return True

    def is_connected(self, x, y):
        return self.find(x) == self.find(y)

    def get_root_part(self):
        part = defaultdict(list)
        n = len(self.root)
        for i in range(n):
            part[self.find(i)].append(i)
        return part

    def get_root_size(self):
        size = defaultdict(int)
        n = len(self.root)
        for i in range(n):
            size[self.find(i)] = self.size[self.find(i)]
        return size

Digit Entry DP

def count(self, num1: str, num2: str, min_sum: int, max_sum: int) -> int:
    MOD = 10 ** 9 + 7
    def f(s: string) -> int:
        @cache
        def f(i: int, sum: int, is_limit: bool) -> int: #sum can be replaced by other state
            if sum > max_sum:
                return 0
            if i == len(s):
                return int(sum >= min_sum)
            ''' if there is no max_sum and min_sum restriction
            if i == len(s):
                return 1    
            '''   
            res = 0
            up = int(s[i]) if is_limit else 9
            for d in range(up + 1):
                res += f(i + 1, sum + d, is_limit and d == up)
            return res % MOD
        return f(0, 0, True)
    ans = f(num2) - f(num1) + (min_sum <= sum(map(int, num1)) <= max_sum)
    return ans % MOD