Zadanie Złote Jabłko 1 etap X OIG

Question

Mam problem z takim zadaniem: https://szkopul.edu.pl/problemset/problem/ZXKiCopRuLe5diKPU0oMNgNn/site/?key=statement

Poza brutem w O(N^2), nie wiem jak to zrobić lepiej, jedynie co to kojarzy mi się trochę to zadanie z zadaniem Megalopolis z 2 etapu XIV OI, żeby puścić DFS-a post-order / pre-order i wyznaczyć dla każdego wierzchołka ile kolejnych jest tego synami i zrobić jakieś drzewo przedziałowe, ale to tu raczej nie zadziała, bo różnych wartości jest nawet 10^5, a nie możemy zrobić tak dużych drzew przedziałowych.

Z góry dziękuję za pomoc i poświęcony czas!

Answer 1

Mam 2 pomysły:

1) w O(nsqrt(n)) - robimy zwykły algorytm Mo. Jeśli dobrze pamiętam jak on działa to będziemy w stanie w czasie sqrt(n) odpowiadać na zapytanie postaci: suma kwadratów liczebności wszystkich kolorów na przedziale [l, r]. Wtedy robimy DFS z pre-order jak w Twoim pomyśle i dla każdego wierzchołka robimy zapytanie dla jego poddrzewa

2) w O(nlog^2n) - dla każdego wierzchołka trzymamy wszystkie kolory, które występują w jego poddrzewie i ich krotności (na jakimś set/ordered map). Ten set będziemy liczyć na podstawie setów dzieci. Jeśli dzieci mają sety s1 i s2, z czego rozmiar s1 < s2 to przenosimy elementy z s1 do s2. Żadnego elementu nie przeniesiemy więcej niż logn razy, a operacje na setach to jakieś logn, więc łącznie O(nlog^2n) jesli dobrze myśle

Comments

A ten drugi pomysł przejdzie czasowo?

---

Trudno powiedzieć. sqrt(n) dla n < 1e5 to mniej więcej to samo co log^2(n)

---

Czyli z tymi setami, to idziemy rekurencyjnie od dołu?

---

Tak

---

ten szkopuł.... przez pół dnia nie działa.....

---

A tak się zastanawiam, napewno wykonamy max N lg(N) operacji przepinania, jak będziemy przepinać mniejszy do większego? Bo faktycznie jak rozrysowywałem przypadki na kartce to wychodzi lekko ponad N, tylko dlaczego N lg N?

---

Ostatnio coś często działa wolno :( Jeśli masz zbiory rozmiaru x1 i x2 (x1 <= x2) i przenosisz elementy z x1 do x2 to rozmiar powstałego zbioru wynosi x1 + x2 >= 2*x1. Z czego wynika że gdy przenosisz jakiś element to on trafia do zbioru co najmniej 2 razy większego. Czyli nie przeniesiesz go więcej niż logn razy

---

No wsm racja, fajny dowód.

---

Przeszło to rozwiązanie z przepinaniem mniejszego do większego na 100pkt, kod:

#include <iostream>
#include <vector>
#include <algorithm>
#include <unordered_map>
#include <queue>

using namespace std;
typedef long long ll;

int n = 0, a = 0, b = 0, nr = 0;
vector<int> A;
vector<vector<int>> krawedzie;
vector<int> porzadek_topologiczny;
vector<bool> czy_bylismy;
queue<int> Q;
vector<unordered_map<int,int>> stat;
vector<ll> wyn;
vector<int> idx_mapy;

int main()
{
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);

    cin >> n;

    krawedzie.assign(n+1,{});
    A.assign(n+1,0);
    for (int i = 1; i <= n; ++i)
        cin >> A[i];
    for (int i = 0; i < n-1; ++i)
    {
        cin >> a >> b;
        krawedzie[a].push_back(b);
        krawedzie[b].push_back(a);
    }

    if (n == 1)
    {
        cout << A[1] << '\n';
        return 0;
    }

    czy_bylismy.assign(n+1,false);
    czy_bylismy[1] = true;
    Q.push(1);
    porzadek_topologiczny.push_back(1);
    while(!Q.empty())
    {
        int spr = Q.front();
        Q.pop();
        for (int i = 0; i < krawedzie[spr].size(); ++i)
        {
            int wierz = krawedzie[spr][i];
            if (czy_bylismy[wierz] == false)
            {
                Q.push(wierz);
                porzadek_topologiczny.push_back(wierz);
                czy_bylismy[wierz] = true;
            }
        }
    }

    idx_mapy.assign(n+1,-1);
    wyn.assign(n+1,0);
    for (int i = n-1; i >= 0; --i)
    {
        int spr = porzadek_topologiczny[i];
        if (spr != 1 and krawedzie[spr].size() == 1)
        {
            wyn[spr] = 1;
            idx_mapy[spr] = stat.size();
            stat.push_back({{A[spr],1}});
        }
        else
        {
            int max_liczebnosc = 0, idx_max = -1, odwolanie = -1;
            ll val = 0;
            for (int j = 0; j < krawedzie[spr].size(); ++j)
            {
                int wierz = krawedzie[spr][j];
                if (idx_mapy[wierz] == -1)
                    continue;
                if (stat[idx_mapy[wierz]].size() > max_liczebnosc)
                {
                    max_liczebnosc = stat[idx_mapy[wierz]].size();
                    idx_max = wierz;
                }
            }
            odwolanie = idx_mapy[idx_max];
            val = wyn[idx_max];
            auto it = stat[odwolanie].find(A[spr]);
            if (it != stat[odwolanie].end())
            {
                val -= (ll)stat[odwolanie][A[spr]] * (ll)stat[odwolanie][A[spr]];
                stat[odwolanie][A[spr]]++;
                val += (ll)stat[odwolanie][A[spr]] * (ll)stat[odwolanie][A[spr]];
            }
            else
            {
                stat[odwolanie][A[spr]] = 1;
                val++;
            }
            for (int j = 0; j < krawedzie[spr].size(); ++j)
            {
                int wierz = krawedzie[spr][j];
                if (idx_mapy[wierz] == -1 or wierz == idx_max)
                    continue;
                int idx = idx_mapy[wierz];
                for (auto it = stat[idx].begin(); it != stat[idx].end(); ++it)
                {
                    int ile = it->second, kolor = it->first;
                    auto it2 = stat[odwolanie].find(kolor);
                    if (it2 != stat[odwolanie].end())
                    {
                        val -= (ll)stat[odwolanie][kolor] * (ll)stat[odwolanie][kolor];
                        stat[odwolanie][kolor] += ile;
                        val += (ll)stat[odwolanie][kolor] * (ll)stat[odwolanie][kolor];
                    }
                    else
                    {
                        stat[odwolanie][kolor] = ile;
                        val += (ll)ile * (ll)ile;
                    }
                }
            }
            wyn[spr] = val;
            idx_mapy[spr] = odwolanie;
        }
    }

    for (int i = 1; i <= n; ++i)
        cout << wyn[i] << ' ';

    return 0;
}

Trzeba trochę uważać, żeby nie kopiować większego, czyli trzeba trzymać dla każdego wierzchołka jaki idx seta go reprezentuje. Puszczamy się albo rekurencyjnie(ja puściłem BFS-a od korzenia(jedynki), żeby było łatwiej.). Powiedziałbym jeszcze jedną rzecz. To rozwiązania ma jak wspomniałeś O(N * lg(N) * lg(N)), ale tak w praktyce powinno często śmigać szybko, bo często kolory przepinamy razem. 

A co do algorytmu Mo, to już muszę się go koniecznie nauczyć, bo to już kolejne zadanie na niego...

Ale ten pomysł z przepinaniem mniejszego do większego, chyba najłatwiejszy. I to spostrzerzenie, że nie przepniemy więcej niż N lg N jest ekstra!

Dzięki!

---

I warto wspomnieć, że gdyby nie chodziło o sumę kwadratów cyfr, tylko poprostu sumę, to te drzewo przedziałowe + DFS preorder, powinien zadziałać w O(N lg N).

Answer 2

Już sobie poradziłem, faktycznie to drzewo przedziałowe wchodzi. Puszczam DFS-a preorder / postorder od korzenia, dla każdego wierzchołka naliczam ile jest wierzchołków w poddrzewie ukorzenionym w nim, i porządek przeglądania tym DFS-em. Robię tablicę vectorów, idx-y wierzchołków o danej wartości, sortuje je własnym operatorem(po tym porządku jak przechodziłem DFS-em) i robię zwykłe drzewo przedziałowe, przeglądam kolor po kolorze, tylko, żeby złożonność się nie zkwadraciła, to muszę jak przeglądne dany kolor to w drzewie przedziałowym nie mogę go clearować całego drzewa przedziałowego, tylko muszę odznaczyć te elementy co zaznaczyłem. Dzięki temu zaznaczę N elementów, i odznaczę N.

O(N lg N), ze względu na sortowanie i operacje na drzewie przedziałowym punkt-przedział.

Comments

Nie widzę jak to zrobić jakimkolwiek drzewem przedziałowym

---

A masz jakiś pomysł czy da się to zrobić tym sposobem z naliczeniem tych statystyk? Jakiś pierwiastek albo coś takiego?

---

@pasjonat_algorytmiki, 

Podsumuwując ten pomysł jest zły.

---

O proszę proszę, robiąc zadanie z elompu z contestu o drzewach i DFS-ie natknąłem się na zadanie, które można zrobić dokładnie tym sposobem z drzewem punkt-przedział i DFS-em preorder, czyli mimo tego, że w tym zadaniu było to nie przydatne, to oddało w innym ;)

Link do treści: https://www.eolymp.com/en/contests/30114/problems/351155

Btw eolymp, szkopuł, sio2, codeforces, usaco i cses, to chyba moim zdaniem najlepsze platformy. Na eolymp jest dużo zadań z EJOI, olimpiad z innych krajów np. Ukrainy czy Azerbejdżanu, zadania z USACO i też dużo contestów tematycznych, z zazwyczaj trudnymi zadaniami, np na drzewa,sortowanie topologiczne, dynamiki N^3, dynamiki liniowe.....

---

A co gdyby kolory można było zmieniać? Jak to wtedy rozwiazać? Btw to jest całkiem trudne :)