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

Bisect:

index1 = bisect.bisect_left(num, val)  #lower_bound
index2 = bisect.bisect_right(num, val)  #upper_bound

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 = {}

Sorted Containers

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. Here we use a library Sorted Containers to address thus. Compared with the data structures implemented by red and black trees in other languages, its unique advantage is the ability to perform random access in $O(logn)$ time.

More usages are at sortedcontainers.

from sortedcontainers import SortedList
sl = SortedList('abbcccddddeeeee')
sl.add('f')
sl.remove('f')
c_pos = sl.index('c')
third_element = sl[3]
sl.bisect_left('d') # 6
sl.bisect_right('d') # 10
 
from sortedcontainers import SortedDict
sd = SortedDict({'b': 2, 'c': 3, 'd': 4, 'e': 5})
# same usage as dict and sorted list

from sortedcontainers import SortedSet
# use as a unique-element sorted list and a set

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)
heapq.heapreplace(nums, nums[0]*2)

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).
4. custom sort:

class Wapper:
    def __init__(self, x: int, y: int) -> None:
        self.x = x
        self.y = y

    def __lt__(self, other: "Wapper") -> bool: # overload <
        return self.x * other.y < self.y * other.x

5. Reduce: Apply a function of two arguments cumulatively to the items of a sequence or iterable, from left to right, so as to reduce the iterable to a single value.
   We can use it with functions in operator.

import operator
from functools import reduce
res = reduce(iand, [1, 2, 3, 4, 5]) # res = 1&2&3&4&5

Some useful operators (https://docs.python.org/3/library/operator.html):

Operation	Syntax	Function
Addition	a + b	add(a, b)
Concatenation	seq1 + seq2	concat(seq1, seq2)
Division	a / b	truediv(a, b)
Division	a // b	floordiv(a, b)
Bitwise And	a & b	and_(a, b)
Bitwise Exclusive Or	a ^ b	xor(a, b)
Bitwise Or	a | b	or_(a, b)
Multiplication	a * b	mul(a, b)

6. Tips in bit operations:

m = x & (x-1) # drop last 1 in x
n = x & (-x)  # get last 1 in x

Be aware the potential overflow problem in -x if using C++.

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

    def get_root_num(self):
        n = len(self.root)
        return len([i for i in range(n) if self.root[i]==i])

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

Fenwick Tree

Be aware that it’s 1-index.

class BIT:
    def __init__(self, n):
        self.tree = [-inf] * (n+1) # pre max
        self.tree = [0] * (n+1)	   # pre sum

    def update_max(self, i: int, val: int) -> None:
        while i < len(self.tree):
            self.tree[i] = max(self.tree[i], val)
            i += i & -i

    def pre_max(self, i: int) -> int:
        mx = -inf
        while i > 0:
            mx = max(mx, self.tree[i])
            i &= i - 1
        return mx

    def update_sum(self, i: int, val: int) -> None:
        while i < len(self.tree):
            self.tree[i] += val
            i += i & -i

    def pre_sum(self, i: int) -> int:
        sm = 0
        while i > 0:
            sm += self.tree[i]
            i &= i - 1
        return sm