7. Organizácia projektov vzorové riešenie

5 2 2
1 3 4 5 2
5 3 2 1 4

18

4 2 2
10 8 8 3
10 7 9 4

31

5 3 1
5 2 5 1 7
6 3 1 6 3

23

Táto úloha sa dala riešiť viacerými spôsobmi. V tomto vzorovom riešení si postupne ukážeme tri rôzne prístupy, ktoré sa dali na riešenie použiť. Napriek tomu, že len jeden je dostatočne rýchly, každý z nich prináša iný pohľad na ten istý problém.

Dynamické programovanie

V našej úlohe sa vlastne snažíme rozdeliť ľudí do troch skupín – tím \(A\), tím \(B\) a skupinu, ktorú musíme prepustiť. Pri dynamickom programovaní sa snažíme rozdeliť problém na menšie podproblémy. Môžeme sa inšpirovať problémom batohu¹, kde sme sa snažili rozdeliť veci na dve kopy – veci, ktoré dáme do batoha a veci, ktoré do batoha nedáme. Podproblémom bolo, že chceme nájsť optimálne riešenie pre prvých \(i\) vecí a kapacitu batoha \(j\).

V našom prípade by sme si mohli povedať, že chceme nájsť optimálne riešenie pre prvých \(i\) programátorov, pričom v tíme \(A\) sa z nich bude nachádzať \(j\) programátorov a v tíme \(B\) sa bude nachádzať \(k\) programátorov. Označme si maximálny súčet skúseností oboch tímov pre takýto podproblém ako \(P[i][j][k]\). Náš podproblém je jednoznačne určený trojicou čísel \((i,j,k)\).

Už sme si definovali, čo je náš podproblém a teraz nám už ostáva iba sa zamyslieť, ako vieme vypočítať nové riešenie z prechádzajúcich hodnôt. Zoberme si \(i\)-teho programátora. Kde sa môže tento programátor nachádzať v optimálnom riešení? Samozrejme, máme tri možnosti – buď je v tíme \(A\), v tíme \(B\), alebo sa nenachádza v žiadnom tíme. Rozoberme si všetky tieto možnosti.

Nech sa nachádza v tíme \(A\). Potom odobratím tohto programátora z tímu \(A\) získame optimálne riešenie pre podproblém \((i-1, j-1, k)\). Prečo? Môžeme to dokázať sporom. Nech hodnota tohto riešenia je \(r\). Predpokladajme, že toto nie je optimálne riešenie pre \((i-1, j-1, k)\), čiže \(r < P[i-1][j-1][k]\). Potom vieme zobrať optimálne riešenie pre \((i-1, j-1, k)\) a pridať \(i\)-teho programátora do tímu \(A\), čím dostaneme lepšie riešenie ako \(P[i][j][k]\), lebo \(P[i][j][k] = r + a_i < P[i-1][j-1][k] + a_i\). Tým sme sa však dostali do sporu. Čiže odobratím \(i\)-teho programátora sme museli nutne dostať riešenie s hodnotou \(P[i-1][j-1][k]\). Tým pádom vieme povedať, že ak programátor \(i\) skončí v tíme \(A\), tak \(P[i][j][k] = P[i-1][j-1][k] + a_i\).

Čo ak sa programátor \(i\) nachádza v optimálnom riešení pre \((i, j, k)\) v tíme \(B\)? Potom ak ho odoberieme z tímu \(B\), tak dostaneme optimálne riešenie pre podproblém \((i-1, j, k-1)\). Dôkaz je znova podobný tomu predchádzajúcemu. V takomto prípade bude platiť, že \(P[i][j][k] = P[i-1][j][k-1] + b_i\).

Posledná možnosť je, že \(i\)-ty programátor sa nenachádza v žiadnom tíme. Potom platí,že toto riešenie je optimálnym riešením aj pre \((i-1, j, k)\), čiže \(P[i][j][k] = P[i-1][j][k]\).

Ak si to zhrnieme, tak potom \(P[i][j][k]\) vieme vypočítať ako maximum z troch hodnôt: \(P[i-1][j-1][k] + a_i\), \(P[i-1][j][k-1] + b_i\) a \(P[i-1][j][k]\).

Ešte si musíme vyjasniť, ako inicializujeme pole P[][][]. Triviálnym prípadom je podproblém, keď máme \(0\) programátorov. V takom prípade vieme inicializovať \(P[0][0][0] = 0\). Ostatné hodnoty \(P[0][j][k]\) inicializujeme na mínus nekonečno, lebo tieto prípady nemajú riešenie. Ak totiž máme nula programátorov, tak neviem mať v žiadnom tíme nenulový počet programátorov.

Časová zložitosť tohto riešenia je \(O(n \cdot x \cdot y)\). Pamäťová zložitosť je tiež \(O(n \cdot x \cdot y)\), ale dá sa zlepšiť, ak si všimneme, že na výpočet hodnoty \(P[i][j][k]\) potrebujeme iba niekoľko políčok okolo a zvyšné môžeme zabudnúť.

Párovanie a toky

Ďalší pohľad je grafový, využívajúci niekoľko pomerne štandardných algoritmov. Keďže toto riešenie stále nie je vzorové, tak tieto algoritmy nevysvetľujeme do podrobnosti. Ak ich teda nepoznáte, nič si z toho nerobte. Vo vzorovom riešení ich vôbec používať nebudeme.

Tento problém sa dá preformulovať aj ako problém maximálneho váhovaného párovania. Zostrojme si bipartitný graf. Vrcholy v prvej partícii zodpovedajú programátorom a vrcholy v druhej partícii zodpovedajú pozíciám v tíme. Teda prvá partícia má \(n\) vrcholov a druhá \(x+y\) vrcholov. Medzi každými dvoma vrcholmi z rôznych partícií vedie hrana, pričom ak máme hranu z \(i\)-teho programátora do \(j\)-tej pozície v tíme \(A\), tak cena tejto hrany je \(a_i\) a ak máme hranu do tímu \(B\), tak cena tejto hrany je \(b_i\). Na takýto graf potom vieme použiť niektorý z algoritmov na hľadanie maximálneho váhovaného párovania na bipartitnom grafe. Ani jeden však nebude dostatočne rýchly, keďže náš graf je pomerne veľký a hlavne obsahuje až \(n(x+y)\) hrán.

Problém maximálneho párovania sa však dá preformulovať na problém maximálneho toku. Stačí nám pridať dva špeciálne vrcholy. Prvý vrchol bude pospájaný so všetkými vrcholmi v prvej partícii a druhý vrchol bude pospájaný so všetkými vrcholmi v druhej partícii. Prvý vrchol je zdroj (source) a druhý je odtok (sink). Kapacita každej hrany je jedna.

Môžeme si všimnúť, že v druhej partícii máme zbytočne veľa vrcholov. Všetky pozície v tíme \(A\) predsa vieme skomprimovať do jedného vrchola a z tohto vrchola pridať hranu do odtoku s kapacitou \(x\). A rovnako pre vrcholy patriace tímu \(B\).

Na tomto grafe potom môžeme spustiť nejaký všeobecný algoritmus na hľadanie maximálneho toku s maximálnou cenou². Zmenšením druhej partície z \(x+y\) na \(2\) sme zmenšili počet hrán medzi týmito partíciami z \((x+y)n\) na \(2n\), čím dostaneme lepšiu časovú zložitosť. Bohužiaľ, ani toto nestačí na vzorové riešenie.

Vzorové riešenie

Vzorové riešenie je v podstate greedy algoritmus. Doteraz sme riešili o dosť všeobecnejšie problémy. Teraz budeme postupovať trocha ináč. Pozrieme sa na to, ako fungujú niektoré špeciálne prípady nášho problému.

Zamyslime sa nad prípadom, keď bude \(y=0\). V tomto prípade sa nám oplatí utriediť programátorov podľa hodnoty \(a_i\) zostupne a zobrať prvých \(x\) programátorov.

Čo ak sa \(y=1\)? Nech máme znova utriedených programátorov podľa \(a_i\). Potom sa nám môže oplatiť zobrať jedného z prvých \(x\) programátorov a dať ich do tímu \(B\) a doplniť programátora číslo \(x+1\) do \(A\). V opačnom prípade zoberieme prvých \(x\) programátorov do tímu \(A\) a zo zvyšku zoberieme programátora s najvyšším \(b_i\) a dáme ho do tímu \(B\).

Čo ak sa \(y=2\)? Určite vieme povedať, že prvých \(x\) programátorov bude v tíme \(A\) alebo \(B\). Nechceme ich teda prepustiť. Vyberme prvého programátora do tímu \(B\). Môže sa nám oplatiť zobrať niektorého z prvých \(x\) programátorov (opäť usporiadaných podľa \(a_i\)). V takom prípade potrebujeme doplniť tím \(A\), čo samozrejme spravíme \((x+1)\)-vým programátorom. Ostáva ešte druhý človek do tímu \(B\). A opäť to môže byť niekto z tímu \(A\), alebo niekto úplne mimo. Ak je to niekto z \(A\), tak tento tím doplníme \((x+2)\)-hým programátorom.

Všimnime si nasledovný fakt: Ak si usporiadame všetkým programátorov podľa čísla \(a_i\), tak v každom optimálnom riešení existuje taká hodnota \(k\), že všetci programátori z tímu \(A\) sú medzi prvými \(k\) programátormi. Naviac vieme, že ak máme najmenšie také \(k\), tak presne \(k-x\) z týchto \(k\) programátorov musíme umiestniť do tímu \(B\) a zo zvyšných \(n-k\) programátorov musíme umiestniť \(y-(k-x)\) programátorov do tímu \(B\). Otázkou ostáva, či pre zadané \(k\) vieme efektívne vypočítať, ktorých programátorov kam umiestniť.

Z posledných \(n-k\) programátorov chceme do tímu \(B\) vybrať \(y-(k-x)\) tých, ktorý majú najväčšiu hodnotu \(b_i\). Musíme sa teda ešte zamyslieť, ako rozdeliť do tímov prvých \(k\) programátorov. Predstavme si, že sme všetkých \(k\) programátorov dali do tímu \(A\). Čo sa stane, ak \(i\)-teho z nich presunieme do tímu \(B\)? Z tímu \(A\) stratíme \(a_i\) skúseností a do tímu \(B\) pribudne \(b_i\) skúseností. Takže celkový zisk/strata bude \(b_i-a_i\). No a zjavne chceme presunúť tých \(k-x\) programátorov, pri ktorých získame čo najviac, teda pre ktorých bude číslo \(b_i-a_i\) čo najväčšie.

Nájsť optimálne riešenie teda vieme nasledovne: Programátorov si zoradíme podľa hodnoty \(a_i\). Potom vyskúšame každú prípustnú hodnotu \(k\), teda \(x \leq k \leq x+y\). Pre dané \(k\) vyberieme z prvých \(k\) programátorov \(k-x\) tých, ktorý majú najväčšiu hodnotu \(b_i-a_i\) a týchto programátorov dáme do tímu \(B\). Zvyšných z prvej \(k\)-tice dáme do tímu \(A\). Následne z ostatných programátorov vyberieme \(y-(k-x)\) programátorov s najvyšším \(b_i\), ktorých zaradíme do tímu \(B\). Pre každú hodnotu \(k\) dostaneme nejaké riešeníe a optimálne bude to najlepšie z nich.

Ostáva už len navrhnúť rýchly algoritmus na zostrojenie tohto riešenia. V prvej časti spočítame pre každé \(k\), koľko najviac vieme získať ak zoberieme prvých \(k\) do tímu \(A\) a následne \(k-x\) z nich hodíme do tímu \(B\). V druhej časti spočítame pre každé \(k\) ako vieme najlepšie nahrabať zvyšných \(y-(k-x)\) programátorov do tímu \(B\). Aké dátové štruktúry budeme potrebovať? Aké operácie budeme často vykonávať? Pre oba problémy sa nám bude hodiť dátová štruktúra, ktorá dovolí efektívne vkladať hodnoty a vyberať najväčší prvok – teda halda.

V prvej časti utriedime programátorov podľa \(a_i\). Potom prechádzame cez takto utriedených programátorov a rátame si súčet \(a_i\). V halde si udržiavame pre každého programátora v tíme \(A\) držať rozdiel \(b_i-a_i\). Chceme v nej teda mať najviac \(x\) prvkov.

Postupne prechádzame programátorov, začínajúc od tých s najvyšším \(a_i\). V každej iterácii priradíme ďalšieho programátor do tímu \(A\) a jeho rozdiel vložíme do haldy. Ak máme v halde viac ako \(x\) prvkov, tak vyberieme von programátora s najväčším rozdielom a preradíme ho do tímu \(B\). Počítame si súčet rozdielov \(b_i-a_i\) pre programátorov, ktorých sme vybrali z haldy a tiež si počítame súčet \(\sum_{i=1}^{k} a_i\). Súčet rozdielov pre nejaké \(k\) je náš zisk, ktorý dostaneme ak daných programátorov presunieme do tímu \(B\). K tomuto číslu ešte môžeme pričítať \(\sum_{i=1}^{k} a_i\). Toto dokopy tvorí súčet skúseností tímu \(A\) a tímu \(B\), ktorý sme doteraz dosiahli pre dané \(k\).

V tomto čísle sa nenachádza súčet zvyšných \(y-(k-x)\) programátorov, ktorých musím pridať do tímu \(B\). Toto spočítame v druhej fáze. Aby sme tieto súčty spočítali, tak budeme iterovať cez programátorov odzadu v poradí v ako sa nachádzajú v utriedenom poli hodnôt \(a_i\). Iterujeme od \(i=n\) až po \(i=x\). Tentokrát si v halde udržiavame hodnoty \(b_i\). V každej iterácii pridáme novú hodnotu \(b_i\). Ak \(i < x+y\), tak vyberieme jedného programátora s najväčším \(b_j\). Počítame si súčet \(b_j\) pre vybratých programátorov. Všimnime si, že pre \(k=x+y\) nevyberáme ešte žiadneho programátora. Pre \(k=x+y-1\) však už potrebujeme do tímu \(B\) doplniť jedného programátora a preto vyberieme toho s najväčším \(b_i\). Súčet hodnôt \(b_i\) vybranných programátorov korešponduje súčtu skúseností zvyšných \(y-(k-x)\) programátorov, ktorých pridáme do tímu \(B\).

Zhrňme si to na záver. V prvej fáze algoritmu sme počítali koľko skúseností získame ak zoberieme prvých \(k\) programátorov a z nich \(k-x\) s najväčším rozdielom \(b_i-a_i\) dáme do \(B\) a zvyšných \(x\) dáme do \(A\). V druhej fázi sme pre každé \(k\) počítali, koľko skúseností vieme ešte navyše získať ak doplníme \(y-(k-x)\) programátorov do tímu \(B\). Ak sčítame súčet, ktorý sme dostali pre nejaké \(k\) v prvej fáze sčítame s číslom, ktoré sme dostali v druhej fáze pre to isté \(k\), tak dostaneme celkový najväčší súčet skúseností, ktorý vieme dostať pri rozdelení definovanom hodnotou \(k\). Z týchto všetkých \(k\) nám už ostáva iba vybrať to najlepšie.

Celková časová zložitosť nášho algoritmu je \(O(n \log n)\), lebo sme potrebovali triediť a používali sme haldu. Pamäťová zložitosť je \(O(n)\).

from itertools import accumulate
from heapq import heappop, heappush


def top(ppl_indices, vals, start):
    """ Pre kazde i>=start vypocitaj sucet top (i-start) hodnot z pola vals.
    Pole vals iteruj v poradi urcenim polom ppl_indices."""
    queue = []
    res = [0 for i in range(len(ppl_indices))]
    for k, idx in enumerate(ppl_indices):
        # Pythonova halda je MIN-halda a my chceme MAX-haldu.
        # Preto vkladame zaporne hodnoty do haldy.
        heappush(queue, -vals[idx])
        if k >= start:
            res[k] = res[k-1] - heappop(queue)

    return res


n, a_size, b_size = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))

conversion_gain = [y - x for x, y in zip(a, b)]
ordered_by_a = sorted(zip(a, range(n)), reverse=True)
prefix_sums_a = list(accumulate([x for x, _ in ordered_by_a]))
# Kolko ziskame konverziou z timu A to timu B?
conversions = top([idx for _, idx in ordered_by_a], conversion_gain, a_size)
# Dopln zvysnych do timu B.
rest_of_bs = list(reversed(top([idx for _, idx
                                in reversed(ordered_by_a[a_size:])],
                               b, n - a_size - b_size))) + [0]

# Scitaj vsetko dokopy.
sol = max(prefix_a + convert + add_bs
          for prefix_a, convert, add_bs
          in zip(prefix_sums_a[a_size-1:a_size+b_size],
                 conversions[a_size-1:a_size+b_size],
                 rest_of_bs))

print(sol)