stoer-wagner.py

# Tworzy podzbiory S, T (wierzchołków), których waga minimalnego cięcia jest minimalna
# (czyli suma wag krawędzi, których usunięcie je rozspójni jest najmniejsza)

# na zdjęciu przykład jak to działa

from heapq import heappop, heappush
from checker import check

class heap_tmp:
    def __init__(self, w, v):
        self.w = w
        self.v = v

    def __lt__(self, other):
        return self.w > other.w # i teraz nie trzeba zmieniać wartości na -


class hashed_maxheap:
    def __init__(self):
        self.PQ = []
        self.added_weights = {}

    def push(self, w, v):
        if v in self.added_weights:
            w += self.added_weights[v]

        self.added_weights[v] = w

        heappush(self.PQ, heap_tmp(w, v))

    def pop(self):
        while len(self.PQ) > 0:
            res = heappop(self.PQ)
            w, v = res.w, res.v
            if self.added_weights[v] == w:
                break # rozwiązanie problemu duplikatów w kolejce, lazy deletion

        return w, v
    
    def is_empty(self):
        return len(self.PQ) == 0


class Node:
    def __init__(self):
        self.edges = {}    # słownik  mapujący wierzchołki do których są krawędzie na ich wagi

    def addEdge( self, to, weight):
        self.edges[to] = self.edges.get(to,0) + weight  # dodaj krawędź do zadanego wierzchołka
                                                        # o zadanej wadze; a jeśli taka krawędź
                                                        # istnieje, to dodaj do niej wagę

    def delEdge( self, to ):
        del self.edges[to]                              # usuń krawędź do zadanego wierzchołka


def gen_adj_dict(L):
    n = 0
    for x, y, _ in L:
       x -= 1
       y -= 1
       if x > n: n = x
       if y > n: n = y

    n += 1

    G = {Node(): i for i in range(n)}

    nodes = list(G.keys())

    for x, y, c in L:
        x -= 1
        y -= 1
        nodes[x].addEdge(nodes[y], c)
        nodes[y].addEdge(nodes[x], c)

    return G

# etap 1
# funkcja, która zwraca nam 2 wierzchołki do scalenia, o danej wadze cięcia je rozspójniającego
# realizacja z kopcem, hashmapą (sprawdzamy jej długość i jak == 1 to zwracamy parę (wierzchołek nie w kopcu, pierwszy wierzchołek z kopca))
# z wagą równą sumie wag krawędzi wychodzących z wierzchołka, który nie jest w kopcu (czyli tego co zostanie)

def min_cut_phase(G : dict):
    H = hashed_maxheap()
    nodes = list(G.keys())
    H.push(0, nodes[0])
    # G to lista node'ów
    processed = set()
    last = None
    prev = None

    while not H.is_empty():
        u: Node
        _, u = H.pop()

        if u in processed: continue

        processed.add(u)
        prev = last
        last = u

        for v in u.edges.keys():
            if v not in processed:
                H.push(u.edges[v], v)

    # prev to wierzchołek co był ostatni zdjęty ze stosu
    # last to ten, który pozostał na sam koniec (singleton)

    last : Node
    # od pythona 3.7 kolejność w kluczy w słowniku jest zgodna z ich dodawaniem
    cut_weight = 0
    for w in last.edges.values():
        cut_weight += w

    return prev, last, cut_weight

# etap 2
# łączenie par wierzchołków zwracanych przez etap 1

def merge_vertices(G : dict, x : Node, y : Node):
    for u in y.edges.keys():
        u : Node
        if u != x:
            x.addEdge(u, y.edges[u])
            u.delEdge(y)
            u.addEdge(x, y.edges[u])
    if y in x.edges:
        x.delEdge(y)

    del G[y]


def stoer_wagner(L):
    G = gen_adj_dict(L)
    min_cut = float('inf')

    while len(G) > 1:
        v, u, w = min_cut_phase(G)
        if w < min_cut: min_cut = w
        merge_vertices(G, v, u)
    
    return min_cut


check(stoer_wagner)