Snippets for memorization/contests

template.py

## Delete element from heap in O(logn)
def delete(num, heap):
    index = heap.index(num)

    heap[index] = heap[-1]
    del heap[-1]

    if index < len(heap):
        heapq._siftup(heap, index)
        heapq._siftdown(heap, 0, index)

## GCD
def gcd(x, y):
    if x == 0: return y
    return gcd(y % x, x)

## Fenwick 
# via alex wice
class Fenwick:
    def __init__(self, n):
        sz = 1
        while sz <= n:
            sz *= 2
        self.size = sz
        self.data = [0] * sz

    def sum(self, i):
        s = 0
        while i > 0:
            s += self.data[i]
            i -= i & -i
        return s

    def add(self, i, x):
        while i < self.size:
            self.data[i] += x
            i += i & -i
            
## Merkle tree
def traverse(node):
    if not node: return ''
    s = hash((traverse(node.left), str(node.val), traverse(node.right))
    return s
             
## LIC
def lic(sequence):
    dp = []
    for i in range(len(sequence)):
        i = bisect.bisect_left(dp, nums[i])
        if i == len(dp):
            dp.append(nums[i])
        else:
            dp[i] = nums[i]
    return len(dp)
        
## Dijkstra


## Segment tree
             

## Morris traversal (inorder)
def inorderTraversal(root):
    curr = root
    res = []
    while curr:
        if curr.left:
            pre = curr.left
            while pre.right and pre.right!=curr:
                pre = pre.right
            if pre.right == curr:
                pre.right = None
                res.append(curr.val)
                curr = curr.right
            else:
                pre.right = curr
                curr = curr.left
        else:
            res.append(curr.val)
            curr = curr.right
    return res


# DSU
class DSU:
    def __init__(self, n: int):
      self.target = n - 1
      self.dsu = {i:i for i in range(n)}
      self.counts = [1] * n
      
    def union(self, p, q):
      root1 = self.find(p)
      root2 = self.find(q)
      if root1 == root2: return False
      
      if self.counts[root1] < self.counts[root2]:
        root1, root2 = root2, root1
      
      self.dsu[root2] = root1
      self.counts[root1] += self.counts[root2]
      return True
    
    def find(self, p):
      while p != self.dsu[p]:
        self.dsu[p] = self.dsu[self.dsu[p]]
        p = self.dsu[p]
      return p
    
    @property
    def connected(self):
      return self.find(0) == self.find(self.target)
    
    def size(self, p):
        return self.counts[self.find(p)]
             
## Kruskal's
def MST(heights):
      m, n = len(heights), len(heights[0])
      edges = self.gather_edges(heights)
      heapq.heapify(edges)
      DSU = DSU(m * n)
      v = 0 
      
      while not DSU.connected:
        v, root, child = heapq.heappop(edges)
        DSU.union(root, child)
      
      return v 
    
    def gather_edges(self, D):
      edges = []
      m,n = len(D), len(D[0])
      
      for i,j in itertools.product(range(m), range(n)):
        from_id = n * i + j
        if i:
          value = abs(D[i][j] - D[i - 1][j])
          to_id = from_id - n
          edges.append([value, from_id, to_id])
        
        if j:
          value = abs(D[i][j] - D[i][j - 1])
          to_id = from_id - 1
          edges.append([value, from_id, to_id])
          
      return edges
 
## Prim's
def MST(A):
  graph = defaultdict(list)
  q = [(0, 0)]
  visited = set()

  for i, (x1,y1) in enumerate(A):
    for j, (x2, y2) in enumerate(A[i + 1:], i + 1):
      d = abs(x2 - x1) + abs(y2 - y1)
      graph[i].append((d, j))
      graph[j].append((d, i))

  cost = 0
  while len(visited) < len(A):
    d, i = heapq.heappop(q)
    if i not in visited:
      visited.add(i)
      cost += d

      for t in graph[i]:
        heapq.heappush(q, t)

  return cost

## Rabin-Karp            
def findString(haystack: str, needle: str) -> int:
    n, h = len(needle), len(haystack)
    if n > h: return -1

    # multiplier, based on base digits/letters
    a = 26 #10
    mod = 2**15 #how to choose? by extent of possible values

    needle = [ord(n) - ord('a') for n in needle]
    # this can also be a lambda function if you want to do one pass
    haystack = [ord(h) - ord('a') for h in haystack]

    t = ref_h = 0
    # build the reference hash
    for i in range(n):
        t = (t * a + haystack[i]) % mod
        ref_h = (ref_h * a + needle[i]) % mod

    if t == ref_h: return 0

    # const value to be used often : a**L % modulus
    aL = pow(a, n, mod)

    for start in range(1, h - n + 1):
        # compute rolling hash in O(1) time
        t = (t * a - haystack[start - 1] * aL + haystack[start + n - 1]) % mod
             
        if t == ref_h: return start

    return -1
             
## Stars / bars
n + r - 1 choose r             
             
## Inverse mod
pow(x, mod - 2, mod)    

## Pow   
def pow(x,n,m):
    if n == 0: return 1
    if n == 1: return x
    if n < 0: return pow(1/x,-n,m - 2)
    return (pow(x * x,n // 2, m) % m)* (1 if n % 2 == 0 else x % m)

## LCP with Suffix arrray + Manber Myers
class Solution:
    def build_lcp(self, s, sa):
        n = len(sa)
        rank = [0] * n
        for i in range(n):
            rank[sa[i]] = i

        k = 0
        lcp = [0] * n
        for i in range(n):
            if rank[i] == n - 1:
                k = 0
                continue
            j = sa[rank[i] + 1]
            while i + k < n and j + k < n and s[i + k] == s[j + k]:
                k += 1
            lcp[rank[i]] = k
            k = max(k - 1, 0)
        return lcp

    def suffix_array_manber_myers(self, s):
      result = []
      def sort_bucket(s, bucket, order=1):
          d = defaultdict(list)
          for i in bucket:
              key = s[i:i+order]
              d[key].append(i)
          for k, v in sorted(d.items()):
              if len(v) > 1:
                  sort_bucket(s, v, order*2)
              else:
                  result.append(v[0])
          return result

      return sort_bucket(s, (i for i in range(len(s)))) 
        
    def countDistinct(self, s: str) -> int:
        sa = self.suffix_array_manber_myers(s)
        lcp = self.build_lcp(s, sa)
        n = len(s)
        return ((n * (n + 1)) // 2) - sum(lcp)