Jozwik_solution.py

# O(ElogE + VL^2), gdzie L - liczba lordów
# Indeksowanie od 1 -> więc będę deklarował np tablice o rozmiarze (N + 1)
# Przejść BFSem po grafie, potem dla każdego lorda zrobić zbiory punktów

from data import runtests
from collections import deque

class Node:
    def __init__(self, idx):
        self.idx = idx
        self.out = set()    # zbiór sąsiadów

    def connect_to(self, v):
        self.out.add(v)

# O(V^2)
def lexBFS(G, V):
    sets_list = [set(range(V))]
    ord = []

    while sets_list:
        s = sets_list[-1].pop()
        ord.append(s)
        # teraz trzeba dzielić zbiory na lewe - bez sąsiadów s oraz prawe - z sąsiadami s
        tmp_list = []
        for z in sets_list:
            if not z: continue
            r = z & G[s].out # przecięcie zbiorów - sąsiedzi
            l = z - r # różnica zbiorów - nie-sąsiedzi
            if l:
                tmp_list.append(l)
            if r:
                tmp_list.append(r)

        sets_list = tmp_list
    
    return ord

def kruskal(E, n):
    p = [i for i in range(n)] # parent
    r = [0] * n # rank
    MST = []

    E.sort(key = lambda x:x[2])

    def find(x): # szuka wierzchołka drzewa
        if p[x] != x:
            p[x] = find(p[x])

        return p[x]
    
    def union(x, y):
        x = find(x)
        y = find(y)

        if x == y: return
        if r[x] > r[y]:
            p[y] = x
        else:
            p[x] = y
            if r[x] == r[y]:
                r[y] += 1

    for u, v, w in E:
        if find(u) != find(v):
            union(u, v)
            MST.append((u, v, w))
            if len(MST) == n - 1:
                break

    return MST


def edges_to_list(E, n):
    G = [[] for _ in range(n)]

    for u, v, w in E:
        G[u].append((v, w))
        G[v].append((u, w))

    return G


def BFS(G, s, n, lord): # można go przerwać jak znajdziemy wszystkie zamki lordów
    Q = deque()
    parent = [None] * n # visited zbędne bo mamy drzewo
    cost_sum = [0] * n
    lord_set = set(lord)
    l = len(lord)
    
    Q.append(s)
    cnt = 1

    while Q:
        u = Q.popleft()

        for v, w in G[u]:
            if v != parent[u]: # zamiast visited sprawdzamy tylko krawędź wsteczną
                parent[v] = u
                cost_sum[v] = cost_sum[u] + w
                if v in lord_set:
                    cnt += 1
                    if cnt == l: return parent, cost_sum
                    Q.appendleft(v)
                else:
                    Q.append(v)

    return parent, cost_sum


def convert_to_bitmask(territory):
    bitmask = 0
    for number in territory:
        bitmask |= (1 << number)
    return bitmask


def my_solve(N, streets, lords):
    E = kruskal(streets, N + 1) # MST - posortowane rosnąco po wagach krawędzi
    G = edges_to_list(E, N + 1)

    # terytoria można uzyskać po prostu odpalając bfsa dla każdego lorda od punktu startowego,
    # wygenerowanie terytoriów: O(VL)
    territories = []
    for lord in lords:
        p = set()
        dist = 0
        s = lord[0]
        p.add(s)
        parent, summed_dist = BFS(G, s, N + 1, lord)

        for u in lord:
            start = u
            
            while not u in p:
                p.add(u)
                u = parent[u]

            dist += summed_dist[start] - summed_dist[u]

        territories.append((p, dist))

    # teraz należy skonstruować graf konfliktów, łatwo udowodnić, że będzie on przekątniowy
    # zbudowany zostanie w postaci akceptowanej przez lexBFS

    n = len(lords) # lordów już indeksuję od 0
    L = [Node(i) for i in range(n)] # graf konfliktów (przekątniowy)

    # Konwersja terytoriów na maski bitowe (dla szybszego wykrywania przecięcia)
    bitmasks = [convert_to_bitmask(territory) for territory, _ in territories]

    for i in range(n):
        for j in range(i + 1, n):
            if bitmasks[i] & bitmasks[j]:  # użycie operacji bitowej AND
                L[i].connect_to(j)
                L[j].connect_to(i)

    # algorytm na wyznaczenie maksymalnego niezależniego zbioru z pracy naukowej FrankP1
    dist = [territories[i][1] for i in range(n)]
    colours = [None] * n # 0 - czerwony, 1 - niebieski
    order = lexBFS(L, n)
    reversed_order = order[::-1]
    og_idx_to_reversed_order = [None] * n
    for i in range(n):
        og_idx_to_reversed_order[reversed_order[i]] = i

    for i in reversed_order:
        if dist[i] > 0:
            colours[i] = 0
            for j in L[i].out:
                if og_idx_to_reversed_order[j] > og_idx_to_reversed_order[i]:
                    dist[j] -= dist[i]
                    if dist[j] < 0: dist[j] = 0 # to zbędne
            dist[i] = 0 # to chyba też

    for i in order:
        if colours[i] == 0 and not any(colours[j] == 1 for j in L[i].out):
            colours[i] = 1

    max_dist = 0
    for i in range(n):
        if colours[i]:
            max_dist += territories[i][1]

    return max_dist


runtests(my_solve)