Union-find в Python

Проблем

Един клиент, три профила

Магазинът има 200 000 клиентски профила. Една и съща клиентка се е регистрирала три пъти: веднъж с телефон, веднъж с имейла от работата, веднъж през приложението. Получава три бюлетина, точките ѝ за лоялност са разделени на три, а отчетът казва, че имаш повече клиенти, отколкото имаш.

Правилото е просто: два профила с общ имейл или общ телефон са един клиент. Но то е заразно. Ако профил a има общ телефон с b, а b общ имейл с c, тогава a, b и c са един клиент, макар че a и c нямат нищо общо.

Как да намериш всички групи, без да сравняваш всеки профил с всеки?

Идеята

Първата идея: етикет на групата

Всеки профил носи етикет: номера на групата си. Отначало всеки е своя група. Когато два профила се окажат свързани, цялата група на втория получава етикета на първия.

class Relabel:
    def __init__(self, n):
        self._label = list(range(n))
        self.count = n

    def find(self, x):
        return self._label[x]

    def union(self, a, b):
        old, new = self._label[b], self._label[a]
        if old == new:
            return False
        for i, group in enumerate(self._label):
            if group == old:
                self._label[i] = new
        self.count -= 1
        return True

Програмата минава през профилите и помни за всеки контакт кой профил го е имал пръв. Ако го среща отново, слива двата профила:

profiles = [
    ["mira@poshta.bg", "0888 100 200"],
    ["0888 100 200", "m.ivanova@firma.bg"],
    ["petar@poshta.bg"],
    ["m.ivanova@firma.bg"],
    ["0899 300 400", "petar@poshta.bg"],
    ["niki@poshta.bg"],
    ["0899 300 400"],
    ["niki@poshta.bg", "0877 500 600"],
    ["mira@poshta.bg", "0899 300 400"],
]
accounts = Relabel(len(profiles))
owner = {}
for p, contacts in enumerate(profiles):
    for contact in contacts:
        if contact in owner:
            accounts.union(owner[contact], p)
        else:
            owner[contact] = p
print("Групи:", accounts.count)
print("Главен:", *(accounts.find(p) for p in range(len(profiles))))

Групи: 2
Главен: 2 2 2 2 2 5 2 5 2

Работи, но всяко сливане минава през всички профили, за да намери кои да преименува. При n профила и около n сливания това са около n² поглеждания: при 200 000 профила десетки милиарди.

Идеята

Гора от родители

Вместо всеки профил да носи етикета на групата, всеки сочи родител: друг профил от групата. Групата е дърво, а профилът, който сочи сам себе си, е коренът: главният профил на групата.

find(x) се качва по родителите, докато стигне до корен. Два профила са в една група, ако find им дава един и същ корен.
union(a, b) намира корените на двата и закача единия под другия. Пипа една клетка от масива, колкото и голяма да е групата.

Сама по себе си гората може да стане дълга верига, и тогава find е бавен. Затова има две подобрения, и заедно те са union-find:

Сливане по размер. Всеки корен знае колко профила има в дървото му. По-малкото дърво отива под по-голямото, затова дълбочината расте бавно.
Компресия на пътя. След като find стигне до корена, минава пътя още веднъж и закача всеки профил по него направо за корена. Следващото find за тях е една стъпка.

Виж го

Сливането отвътре

class UnionFind:
    def __init__(self, n):
        self._parent = list(range(n))
        self._size = [1] * n
        self.count = n

    def find(self, x):
        root = x
        while self._parent[root] != root:
            root = self._parent[root]
        while self._parent[x] != root:
            above = self._parent[x]
            self._parent[x] = root
            x = above
        return root

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self._size[ra] < self._size[rb]:
            ra, rb = rb, ra
        self._parent[rb] = ra
        self._size[ra] += self._size[rb]
        self.count -= 1
        return True

union връща False, ако профилите вече са в една група: тогава нищо не се сменя, и count остава.

Същата програма като горе, с UnionFind вместо Relabel:

profiles = [
    ["mira@poshta.bg", "0888 100 200"],
    ["0888 100 200", "m.ivanova@firma.bg"],
    ["petar@poshta.bg"],
    ["m.ivanova@firma.bg"],
    ["0899 300 400", "petar@poshta.bg"],
    ["niki@poshta.bg"],
    ["0899 300 400"],
    ["niki@poshta.bg", "0877 500 600"],
    ["mira@poshta.bg", "0899 300 400"],
]
accounts = UnionFind(len(profiles))
owner = {}
for p, contacts in enumerate(profiles):
    for contact in contacts:
        if contact in owner:
            accounts.union(owner[contact], p)
        else:
            owner[contact] = p
print("Групи:", accounts.count)
print("Главен:", *(accounts.find(p) for p in range(len(profiles))))

Групи: 2
Главен: 0 0 0 0 0 5 0 5 0

class UnionFind:
    def __init__(self, n):
        self._parent = list(range(n))
        self._size = [1] * n
        self.count = n

    def find(self, x):
        root = x
        while self._parent[root] != root:
            root = self._parent[root]
        while self._parent[x] != root:
            above = self._parent[x]
            self._parent[x] = root
            x = above
        return root

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self._size[ra] < self._size[rb]:
            ra, rb = rb, ra
        self._parent[rb] = ra
        self._size[ra] += self._size[rb]
        self.count -= 1
        return True

Стъпка по стъпка, с думи

Девет клиентски профила, всеки е своя група: _parent[i] = i, всеки профил е корен. Групи: 9.
Профил 1 има телефона на 0: union(0, 1). Корените са различни, 1 се закача под 0. После 3 има имейла на 1: find(1) се качва до 0 и 3 също отива под 0.
Същото за другите: 4 под 2 (общ имейл), 6 под 2 (общ телефон), 7 под 5. Сливането сменя една клетка от масива: родителя на един корен.
Профил 8 има имейла на 0: 8 отива под 0. Групата на 0 вече е с четири профила.
8 има и телефона на 4. find(4) стига до 2, find(8) до 0. Дървото на 2 е по-малко (3 срещу 4), затова коренът му отива под 0: една клетка, три профила сменят групата си.
Накрая програмата вика find за всеки профил. find(4) се качва 4 → 2 → 0, после закача 4 направо за 0: компресия на пътя. Същото за 6.
Две групи, главните профили са 0 и 5. Дървото е плоско: следващото find(4) е една стъпка.

Двете групи са същите като при етикетите, само номерата са други: тук главен е коренът, при етикетите е етикетът, който е оцелял.

Пробвай

Предвиди изхода

Шест профила и шест връзки. След всяко сливане програмата записва броя групи:

uf = UnionFind(6)
links = [[0, 1], [1, 2], [0, 2], [3, 4], [4, 5], [2, 3]]
counts = []
for a, b in links:
    uf.union(a, b)
    counts.append(uf.count)
print(counts)

Отговор

[5, 4, 4, 3, 2, 1]

[1, 1, 1, 1, 1, 1]

[5, 4, 3, 2, 1, 0]

Идеята

Рекурсивният find

Компресията често се пише като рекурсия: find на родителя дава корена, и профилът се закача за него. Всеки профил по пътя прави същото на връщане.

def find(self, x):
    if self._parent[x] != x:
        self._parent[x] = self.find(self._parent[x])
    return self._parent[x]

Пробвай сам:

def find(self, x):
    if self._parent[x] != x:
        self._parent[x] = 
    return self._parent[x]

Със сливане по размер дърветата са ниски: при 200 000 профила най-много 17 нива (log₂ 200 000 ≈ 17,6). Затова рекурсията е безопасна.

По-точно, за любопитните

Без сливане по размер дървото може да стане верига от n профила, и рекурсивният find прави n вложени извиквания, които стоят на стека едно над друго, както в урока „Рекурсия“. Python спира на 1000 вложени извиквания по подразбиране и хвърля RecursionError. Цикълът от урока няма този проблем при никаква форма на дървото.

Идеята

Колко струва

Пусни трите начина. Има n профила и n случайни връзки „тези два имат общ контакт“. Броят се:

при етикетите: всяко поглеждане на етикета на профил при преименуването;
при гората без подобрения и при union-find: всяка стъпка на find, включително погледът към самия корен.

Преименуване на цялата група O(n²)
Гора от родители, без подобрения O(n²)
Union-find: компресия и размер O(n)

Алгоритъм	n = 10	n = 1000	n = 10 000	n = 100 000
Преименуване на цялата група	90	843 000	≈ 84 300 000	≈ 8 430 000 000
Гора от родители, без подобрения	30	25 641	2 102 613	≈ 210 261 300
Union-find: компресия и размер	30	3295	33 355	334 195

Кодът се пуска в твоя браузър на JavaScript. Броят операции е същият на всеки език, защото алгоритъмът е същият. Времето зависи от машината, затова е само за сравнение. „≈“ е оценка по кривата: толкова голямо n не пускаме, за да не замръзне страницата.

class UnionFind:
    def __init__(self, n):
        self._parent = list(range(n))
        self._size = [1] * n
        self.count = n

    def find(self, x):
        root = x
        while self._parent[root] != root:
            root = self._parent[root]
        while self._parent[x] != root:
            above = self._parent[x]
            self._parent[x] = root
            x = above
        return root

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return False
        if self._size[ra] < self._size[rb]:
            ra, rb = rb, ra
        self._parent[rb] = ra
        self._size[ra] += self._size[rb]
        self.count -= 1
        return True

При 1000 профила етикетите правят 843 000 поглеждания, гората без подобрения 25 641 стъпки, а union-find 3 295. При 10 000 гората без подобрения вече прави 2 102 613 стъпки: дърветата ѝ стават дълги вериги. Union-find прави 33 355, около 3,3 стъпки на връзка, и при 100 000 профила 334 195: пак толкова на връзка.

Това е амортизирана цена (виж амортизираната сложност): едно find може да е дълго, но след него пътят е плосък и следващите са кратки. Общо операциите струват почти O(1) всяка.

По-точно, за любопитните

По-точно: със сливане по размер (или по ранг, виж по-долу) и компресия на пътя m операции върху n елемента струват O(m · α(n)). α е обратната функция на Акерман. Тя расте толкова бавно, че за всяко n, което се събира във Вселената, е най-много 4. Затова на практика се казва „почти константа“. Доказателството е на Тарян (1975) и не е кратко.

Вместо размер може да се пази ранг: горна граница на височината на дървото. Коренът с по-малък ранг отива под другия, а рангът расте само при равни рангове. Цената е същата. Размерът е по-удобен, защото отговаря и на въпроса „колко профила има в групата“.

Пробвай

Намери грешката

Колега е написал сливането. Броят групи излиза 2, а главните профили са три различни:

def union(self, a, b):
    ra, rb = self.find(a), self.find(b)
    if ra == rb:
        return False
    if self._size[ra] < self._size[rb]:
        ra, rb = rb, ra
    self._parent[b] = ra
    self._size[ra] += self._size[rb]
    self.count -= 1
    return True

profiles = [
    ["mira@poshta.bg", "0888 100 200"],
    ["0888 100 200", "m.ivanova@firma.bg"],
    ["petar@poshta.bg"],
    ["m.ivanova@firma.bg"],
    ["0899 300 400", "petar@poshta.bg"],
    ["niki@poshta.bg"],
    ["0899 300 400"],
    ["niki@poshta.bg", "0877 500 600"],
    ["mira@poshta.bg", "0899 300 400"],
]
accounts = UnionFind(len(profiles))
owner = {}
for p, contacts in enumerate(profiles):
    for contact in contacts:
        if contact in owner:
            accounts.union(owner[contact], p)
        else:
            owner[contact] = p
print("Групи:", accounts.count)
print("Главен:", *(accounts.find(p) for p in range(len(profiles))))

Групи: 2
Главен: 0 0 2 0 2 5 2 5 0

На кой ред е грешката?

Избери реда с грешката.

Идеята

Къде още

Union-find отговаря на един въпрос: в една група ли са тези два, когато групите само растат. Същото е и при:

зоните на доставка: складове и пътища между тях. Два склада са в една зона, ако има път, може и през други складове;
мрежа от компютри: кои машини се виждат една друга след всеки нов кабел;
алгоритъма на Крускал за най-евтина мрежа от пътища, който ще видиш в модула за графи (algo-graphs). Там union-find казва дали нов път не затваря кръг.

Едно ограничение: union-find не разделя. Ако общ телефон се окаже телефонът на офис и е слял двама различни клиенти, сливането не може да се върне назад. Затова връзките се пазят отделно, а групите се строят наново, когато някоя връзка се махне. Затова и телефон, който се среща в стотици профили, обикновено изобщо не се ползва за сливане.

Упражнение

Зоните на доставка

Реши задачата в редактора си.

Зони. Има 10 склада (0 … 9) и пътища между тях: 0-1, 1-2, 3-4, 5-6, 6-7, 7-5, 8-9. Добави в класа два метода: connected(a, b), който казва дали два склада са в една зона, и size_of(x), който връща колко склада има в зоната на x. Изведи броя зони, дали 0 и 2 и дали 2 и 3 са в една зона, и колко склада има в зоната на 5. После отвори нов път 2-3 и изведи пак броя зони и дали 0 и 4 са в една зона:

Зони: 4
0 и 2: да
2 и 3: не
Зоната на 5: 3 склада
Нов път 2-3. Зони: 3
0 и 4: да

Още задачи

Групите поименно. Изведи всяка група профили като списък: главният профил и всички профили под него.
Най-голямата група. Пази най-големия размер на група, докато сливаш, без да обхождаш всички корени накрая.
Сливане по ранг. Смени размера с ранг и сравни броя стъпки на find с тези от Bench.

Тест

Провери се

1. Профилите a и b имат общ имейл, b и c общ телефон, а d няма общо с никого. Колко клиента са това?4: всеки профил е отделен3: a и c нямат общ контакт2: a, b и c са един клиент, d е друг

2. Какво прави union, когато двата профила вече са в една група?Нищо: find дава един и същ корен, броят групи не се сменяЗакача единия корен под самия негоНамалява броя групи с едно

3. Защо коренът на по-малкото дърво отива под корена на по-голямото, а не обратно?За да остане главен най-старият профилНяма значение, само е по-красивоТака дълбочината расте само при равни размери и не става повече от log₂ n

4. Два профила са слети по общ телефон, а се оказва, че това е телефонът на офис. Как ги разделяш?Махаш реда в масива parentСтроиш групите наново от връзките, без грешнатаС обратен union

Union-find

Бисквитки