Omówienie zadania Bajholmska Firma Transportowa (eliminacje do IOI 2020, dzień 2)

Formalna treść zadania

W zadaniu dane jest drzewo o \(n\) wierzchołkach. Krawędź o numerze \(i\) (dla \(1\leq i < n\)):

• łączy wierzchołki \(u_i\) oraz \(v_i\),
• posiada wagę \(m_i\),
• posiada parametr \(a_i\in\{0,1\}\), który oznacza, że dana krawędź jest ciekawa, jeśli \(a_i=1\) lub że taka nie jest, gdy \(a_i=0\).

Mamy \(q\) zapytań. W \(j\)-tym z nich dana jest liczba \(w_j\). Odpowiedzią na to zapytanie jest liczba ścieżek prostych (czyli takich, w których wierzchołki się nie powtarzają) w drzewie, które spełniają następujące warunki:

1. wagi wszystkich krawędzi na ścieżce są nie mniejsze niż \(w_j\),
2. ścieżka zawiera przynajmniej jedną ciekawą krawędź.

W dalszej części rozwiązania, pisząc ścieżka, będziemy mieć na myśli ścieżki proste.

Co jeśli wagi wszystkich krawędzi są sobie równe? (Podzadanie 2)

Załóżmy, że wagi wszystkich krawędzi są sobie równe. Niech \(m\) będzie ich wspólną wagą. Jeśli dla zapytania o numerze \(j\) mamy \(w_j>m\), to odpowiedź jest oczywiście równa \(0\). Załóżmy dalej, że \(w_j\leq m\). Wówczas możemy zapomnieć o pierwszym warunku z treści zadania, gdyż jest on spełniony dla każdej ścieżki.

Pozostało odpowiedzieć na pytanie: ile jest ścieżek, które zawierają przynajmniej jedną krawędź ciekawą? Z pomocą przychodzi nam technika, która często pojawia się w zadaniach o zliczaniu: zamiast zliczać ścieżki, które spełniają nasz warunek, zliczymy ścieżki, które go nie spełniają.

Dlaczego to nam pomaga? Zauważmy, że jeśli ścieżka nie spełnia warunku, to znaczy, że wszystkie krawędzie na tej ścieżce nie są ciekawe. Okazuje się, że takie ścieżki jest trochę prościej zliczać. Usuńmy z drzewa wszystkie ciekawe krawędzie. Otrzymujemy w ten sposób las złożony z pewnej liczby spójnych drzew. Oznaczmy ich liczbę przez \(k\). Widzimy, że liczba ścieżek w otrzymanym lesie jest równa liczbie ścieżek w oryginalnym drzewie, na których nie ma żadnej ciekawej krawędzi. Jeśli oznaczymy przez \(s_i\) dla \(1\leq i\leq k\) liczbę wierzchołków w spójnym drzewie o numerze \(i\), to liczba ścieżek w naszym lesie będzie równa \[ \sum_{i=1}^{k} \binom{s_i}{2}=\sum_{i=1}^{k} \frac{s_i(s_i-1)}{2}. \] Wiemy, że liczba wszystkich ścieżek w oryginalnym drzewie wynosi \(\binom{n}{2}\), skąd otrzymujemy wzór na liczbę ścieżek, które zawierają przynajmniej jedną krawędź ciekawą: \[ \binom{n}{2}-\sum_{i=1}^{k} \binom{s_i}{2}= \frac{n(n-1)}{2}-\sum_{i=1}^{k} \frac{s_i(s_i-1)}{2}. \]

Liczby \(s_i\) możemy łatwo obliczyć w czasie \(O(n)\), wykonując przeszukiwanie wszerz (znanego także jako DFS) na grafie bez ciekawych krawędzi, skąd otrzymujemy całkowitą złożoność \(O(n)\). Rozwiązanie tego podproblemu pozwala zaliczyć 2 podzadanie. Należy jednak zwrócić uwagę na to, że pomysł, który pojawił się w powyższej analizie - zliczanie ścieżek niespełniających warunku numer 2 - jest kluczowy do pełnego rozwiązania zadania.

Co jeśli w grafie jest tylko jedna ciekawa krawędź? (Podzadanie 3)

Załóżmy, że w drzewie znajduje się tylko jedna ciekawa krawędź. Dla każdego zapytania \(w_j\) zliczamy ścieżki przechodzące przez tę krawędź, na których wagi wszystkich krawędzi wynoszą przynajmniej \(w_j\). Spróbujmy opisać proste rozwiązanie o złożoności \(O(q\cdot n)\). Niech \(u\) i \(v\) będą końcami ciekawej krawędzi. Rozważmy zapytanie o numerze \(j\). Jeśli \(w_j\) jest większe od wagi ciekawej krawędzi, to odpowiedź jest oczywiście równa 0. Załóżmy więc, że \(w_j\) jest mniejsze lub równe tej wadze.

Wykonujemy algorytm przeszukiwania wszerz z wierzchołka \(u\), przechodząc przez krawędzie o wagach nie mniejszych niż \(w_j\) oraz nie przechodząc przez krawędź ciekawą. Analogiczne przeszukiwanie wykonujemy z wierzchołka \(v\). Otrzymujemy w ten sposób dwie liczby \(t_u\) i \(t_v\), które oznaczają liczbę wierzchołków odwiedzonych w przeszukiwaniach rozpoczętych w \(u\) oraz \(v\), odpowiednio. Widzimy, że każda ścieżka interesująca nas w tym zapytaniu musi mieć jeden koniec w pewnym wierzchołku odwiedzonym podczas przeszukiwania z \(u\), a drugi w wierzchołku odwiedzonym podczas przeszukiwania z \(v\). W związku z tym odpowiedź na nasze zapytanie wynosi \(t_u\cdot t_v\). Wartości \(t_u\) i \(t_v\) możemy obliczyć w czasie \(O(n)\) dla danego zapytania, skąd otrzymujemy złożoność \(O(q\cdot n)\).

Spróbujmy przyspieszyć powyższe rozwiązanie. Zauważmy, że fragmenty grafu odwiedzone podczas przeszukiwania wszerz w powyższym pomyśle stają się coraz większe wraz z tym, jak zmniejszamy wartość \(w_j\). Bardziej formalnie, jeśli mamy dwa zapytania o numerach \(j_1\) i \(j_2\) takie, że \(w_{j_1}< w_{j_2}\), to zbiór wierzchołków odwiedzonych przy preszukiwaniu wszerz dla \(j_2\) jest podzbiorem zbioru wierzchołków odwiedzonych przy przeszukiwaniu dla \(j_1\). Nasuwa się więc pomysł, aby posortować zapytania nierosnąco względem wartości \(w_j\) i dynamicznie utrzymywać spójne wierzchołków, które w poprzednim pomyśle zostałyby odwiedzone podczas przeszukiwania wszerz. W ten sposób unikniemy potrzeby wielokrotnego przeszukiwania tych samych wierzchołków.

Aby to sprawnie zrealizować, dla obu spójnych możemy utrzymywać kolejki priorytetowe \(q_u\) i \(q_v\) krawędzi wychodzących z wierzchołków odpowiednich spójnych, które idą do wierzchołków spoza tych spójnych. Na początku do obu kolejek wstawiamy wszystkie krawędzie wychodzące z wierzchołków \(u\) i \(v\), odpowiednio, pomijając krawędź ciekawą. Rozpatrując zapytanie o parametrze \(w_j\), zdejmujemy z kolejki \(q_u\) krawędzie o wagach nie mniejszych niż \(w_j\) i dodajemy do spójnej wierzchołka \(u\) wierzchołki, do których prowadzą te krawędzie. Następnie dodajemy do tej kolejki krawędzie wychodzące z nowo dodanych wierzchołków i powtarzamy ten proces dopóki kolejka nie będzie pusta lub wszystkie krawędzie na kolejce będą miały wagi mniejsze niż \(w_j\). Analogicznie postępujemy z kolejką \(q_v\). Dzięki posortowaniu zapytań otrzymamy poprawne rozmiary spójnych dla każdego z nich. Zauważmy na koniec, że przy jednym zapytaniu możemy wykonać bardzo wiele operacji na kolejkach, jednak sumaryczna liczba operacji wykonanych podczas całego algorytmu będzie \(O(n)\), gdyż w naszym drzewie jest liniowo wiele krawędzi. Otrzymujemy więc złożoność \(O((n+q)\log n)\).

To przyspieszone rozwiązania pozwala zaliczyć 3 podzadanie. Pomysł z sortowaniem zapytań i dynamicznym rozszerzaniem spójnych w drzewie będzie jednak bardzo pomocny w rozwiązaniach kolejnych podzadań.

Rozwiązanie \(O(n^2q)\) (Podzadanie 5)

W tym podzadaniu możemy zliczyć liczbę interesujących nas ścieżek prosto z definicji w czasie \(O(n^2)\) dla każdego zapytania. Ustalmy pewne zapytanie o numerze \(j\). Wybierzmy wierzchołek \(v\) i obliczmy liczbę interesujących nas ścieżek, które zaczynają się w tym wierzchołku. Wykonajmy przeszukiwanie wszerz, zaczynając od \(v\) i przechodząc wyłącznie przez krawędzie o wagach nie mniejszych niż \(w_j\). Dodatkowo, będziemy pamiętać, czy na danej ścieżce podczas przeszukiwania wszerz pojawiła się ciekawa krawędź. Dla każdego odwiedzonego wierzchołka \(u\) należy zwiększyć wynik o 1, jeśli na ścieżce do niego pojawiła się ciekawa krawędź. Powtarzając ten proces dla każdego wierzchołka \(v\), otrzymamy sumaryczną złożoność \(O(n^2)\) na jedno zapytanie, co ostatecznie daje nam złożoność \(O(n^2q)\). To rozwiązanie pozwala zaliczyć podzadanie 5.

Rozwiązanie \(O(n\cdot q)\) (Podzadania 4 i 6)

Załóżmy, że w zadaniu mamy tylko jedno zapytanie. Niech \(w\) będzie wartością parametru z tego zapytania. Wówczas z grafu możemy usunąć wszystkie krawędzie o wagach mniejszych niż \(w\), gdyż zliczane ścieżki nie mogą przechodzić przez takie krawędzie. Otrzymujemy w ten sposób las złożony z pewnej liczby spójnych drzew. Dalej możemy zapomnieć o wagach pozostałych krawędzi, gdyż wszystkie mają wartość co najmniej \(w\). Dla każdego z otrzymanych drzew możemy więc zastosować algorytm opisany wyżej w przypadku podzadania numer 2. Odpowiedzią będzie suma otrzymanych wyników dla każdego z drzew. Otrzymujemy złożoność \(O(n)\) na zapytanie, co powtórzone \(q\) razy daje nam sumaryczną złożoność \(O(n\cdot q)\). Taki algorytm pozwala rozwiązać podzadania 4, 5 i 6.

Rozwiązanie \(O(q\log q + n\log n)\)

W tym rozwiązaniu spróbujemy połączyć pomysł ze zliczaniem ścieżek niespełniających warunku numer 2 z podzadania 2 z pomysłem na sortowanie zapytań z podzadania 3. Przeanalizujmy algorytm opisany w poprzednim punkcie. Dla każdego zapytanie konstruujemy las, odrzucając z naszego drzewa krawędzie o zbyt małych wagach. Następnie pytamy o liczności spójnych w tym lesie oraz o liczności spójnych powstałych na skutek usunięcia wszystkich ciekawych krawędzi z tego lasu.

Spróbujmy zrobić to samo, ale dynamicznie. Posortujmy zapytania nierosnąco względem wartości \(w_j\). Będziemy utrzymywać dwie struktury zbiorów rozłącznych DSU (znanego także jako Find & Union) - jedną dla oryginalnego drzewa, drugą dla drzewa bez ciekawych krawędzi. Oznaczymy je \(D_a\) oraz \(D_b\). Więcej na temat tej struktury można przeczytać w opracowaniu zadania "Małpki" z X OI. Rozważając zapytanie o numerze \(j\) chcemy, aby \(D_a\) zostało zaktualizowane przez wszystkie krawędzie o wagach nie mniejszych niż \(w_j\), a \(D_b\) przez wszystkie nieciekawe krawędzie o wagach nie mniejszych niż \(w_j\). Wówczas odpowiedzią na zapytanie jest różnica liczby ścieżek w \(D_a\) i \(D_b\). Zauważmy, że liczbę różnych ścieżek w lesie utrzymywanym jako struktura zbiorów rozłącznych możemy utrzymywać wraz ze strukturą. Wystarczy zauważyć, że przy łączeniu dwóch spójnych o rozmiarach \(x_1\) i \(x_2\) otrzymujemy dokładnie \(x_1\cdot x_2\) nowych ścieżek. Wystarczy więc dla każdej spójnej trzymać informację o jej rozmiarze i odpowiednio aktualizować wynik przy łączeniu.

Możemy zatem posortować zarówno krawędzie, jak i zapytania nierosnąco względem wag. Będziemy odpowiadać na zapytania w tej kolejności. Gdy odpowiadamy na zapytanie o parametrze \(w_j\), to dodajemy do naszych struktur wszystkie krawędzie o wagach nie mniejszych niż \(w_j\), które jeszcze nie zostały do nich dodane (pamiętając, że do \(D_b\) wrzucamy tylko nieciekawe krawędzie). Następnie odpowiadamy na zapytanie i powtarzamy ten proces dla kolejnego. W ten sposób każda krawędź zostanie dodana do struktury co najwyżej raz, skąd otrzymujemy złożoność \(O(q\log q+n\log n)\), która wynika z konieczności posortowania danych. To rozwiązanie jest wystarczające do tego, aby zaliczyć wszystkie podzadania.