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[^1], 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[^2]. 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)

1. knapsack – https://en.wikipedia.org/wiki/Knapsack_problem↩

2. https://en.wikipedia.org/wiki/Minimum-cost_flow_problem↩