python刷题常用模板

#===================================== 素数筛 Begin =====================================#
MAXN = 1000
prime = []
isprime = [True] * (MAXN + 1)
def euler():
    isprime[1] = False
    for i in range(2, MAXN + 1):
        if isprime[i]: prime.append(i)
        for pr in prime:
            if i * pr > MAXN: continue
            isprime[i * pr] = False
            if i % pr == 0: break
euler()
#===================================== 素数筛 Test =====================================#
assert isprime[1] is False
assert isprime[59] is True
assert isprime[293] is True
print("Euler Prime test passed!")
#===================================== 素数筛 End =====================================#

#===================================== 并查集 Begin =====================================#
class DSU:
    def __init__(self, n):
        self.parent = [i for i in range(n)]

    def find(self, x):
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]
    
    def merge(self, x, y):
        u, v = self.find(x), self.find(y)
        self.parent[u] = v
    
    def same(self, x, y) -> bool:
        return self.find(x) == self.find(y)
#===================================== 并查集 Test =====================================#
dsu = DSU(100)
assert dsu.same(1, 75) is False
dsu.merge(1, 50)
dsu.merge(25, 75)
dsu.merge(50, 25)
assert dsu.same(1, 75) is True
print("DSU test passed!")
#===================================== 并查集 End =====================================#

#===================================== 线段树 Begin =====================================#
class SegmentTree:
    def __init__(self, v):
        self.n = len(v)
        self.n4 = 4 * self.n
        self.tree = [0] * self.n4
        self.lazy = [0] * self.n4
        self.arr = v
        self.root = 1
        self.end = self.n
        self.__build(0, self.end - 1, self.root)
        self.arr = None

    def query(self, l, r):
        return self.__query(1, self.end, self.root, l + 1, r + 1)

    def add(self, l, r, val):
        self.__add(1, self.end, self.root, l + 1, r + 1, val)

    def __build(self, s, t, idx):
        if s == t:
            self.tree[idx] = self.arr[s]
        else:
            m = s + ((t - s) >> 1)
            self.__build(s, m, 2 * idx)
            self.__build(m + 1, t, 2 * idx + 1)
            self.__pushup(idx)

    def __add(self, s, t, idx, l, r, val):
        if l <= s and t <= r:
            self.lazy[idx] += val
            self.tree[idx] += (t - s + 1) * val
        else:
            m = s + ((t - s) >> 1)
            self.__pushdown(s, t, idx)
            if l <= m:
                self.__add(s, m, 2 * idx, l, r, val)
            if m < r:
                self.__add(m + 1, t, 2 * idx + 1, l, r, val)
            self.__pushup(idx)

    def __query(self, s, t, idx, l, r):
        if l <= s and t <= r:
            return self.tree[idx]
        m = s + ((t - s) >> 1)
        self.__pushdown(s, t, idx)
        sm = 0
        if l <= m:
            sm += self.__query(s, m, 2 * idx, l, r)
        if m < r:
            sm += self.__query(m + 1, t, 2 * idx + 1, l, r)
        return sm

    def __pushup(self, idx):
        self.tree[idx] = self.tree[idx * 2] + self.tree[idx * 2 + 1]

    def __pushdown(self, s, t, idx):
        m = s + ((t - s) >> 1)
        if self.lazy[idx] != 0:
            self.lazy[idx * 2] += self.lazy[idx]
            self.tree[idx * 2] += (m - s + 1) * self.lazy[idx]
            self.lazy[idx * 2 + 1] += self.lazy[idx]
            self.tree[idx * 2 + 1] += (t - m) * self.lazy[idx]
            self.lazy[idx] = 0
#===================================== 线段树 Test =====================================#
arr = [1, -2, 3, -4, 5]
seg = SegmentTree(arr)
assert seg.query(0, 2) == 2
seg.add(1, 3, 2)    # arr = [1, 0, 5, -2, 5]
assert seg.query(0, 3) == 4
print("SegmentTree1 test passed!")
#===================================== 线段树 End =====================================#


#===================================== 线段树 （动态开点）Begin =====================================#
class Node:
    def __init__(self, l, r):
        self.left = None
        self.right = None
        self.l = l
        self.r = r
        self.mid = (l + r) >> 1
        self.v = 0
        self.add = 0

class SegmentTree2:
    def __init__(self):
        self.root = Node(1, int(1e9))

    def modify(self, l, r, v, node=None):
        if l > r:
            return
        if node is None:
            node = self.root
        if node.l >= l and node.r <= r:
            node.v = v
            node.add = v
            return
        self.pushdown(node)
        if l <= node.mid:
            self.modify(l, r, v, node.left)
        if r > node.mid:
            self.modify(l, r, v, node.right)
        self.pushup(node)

    def query(self, l, r, node=None):
        if l > r:
            return 0
        if node is None:
            node = self.root
        if node.l >= l and node.r <= r:
            return node.v
        self.pushdown(node)
        v = 0
        if l <= node.mid:
            v = max(v, self.query(l, r, node.left))
        if r > node.mid:
            v = max(v, self.query(l, r, node.right))
        return v

    def pushup(self, node):
        node.v = max(node.left.v, node.right.v)

    def pushdown(self, node):
        if node.left is None:
            node.left = Node(node.l, node.mid)
        if node.right is None:
            node.right = Node(node.mid + 1, node.r)
        if node.add:
            node.left.v = node.add
            node.right.v = node.add
            node.left.add = node.add
            node.right.add = node.add
            node.add = 0
#===================================== 线段树 （动态开点）Test =====================================#
seg = SegmentTree2()
seg.modify(0, 5, 2)
assert seg.query(0, 5) == 2
seg.modify(3, 7, 4)
assert seg.query(5, 5) == 4
print("SegmentTree2 test passed!")
#===================================== 线段树 （动态开点）End =====================================#