4. Súčty v električke vzorové riešenie

Podúloha A

Predstavme si na začiatok, že predtým, ako Kubo nastúpil do električky, v nej ešte neboli žiadni ľudia. Potom by nám stačilo pamätať si, koľko je v električke po ktorej zastávke ľudí. Potom počet ľudí, ktorí nastúpili na nejakom úseku, je jednoducho rozdiel počtu ľudí, ktorí v nej boli predtým, a počtu ľudí, ktorí v nej boli potom.

Toto už je skoro riešenie. Jediný rozdiel je, že na začiatku už nejakí ľudia v električke sú. Nebudeme si teda pamätať ich počet (ten ani nevieme), ale rozdiel oproti počtu na začiatku.

To znamená, že pre každú zastávku si zapamätáme počet ľudí, ktorí až dovtedy nastúpili, čo je súčet počtov pri jednotlivých zastávkach.

Napríklad ak by počty nastupujúcich boli postupne \(2, 3, -7, 5, -4\), potom rozdiely oproti začiatku by boli \(0, 2, 5, -2, 3, -1\). Hodí sa tam mať tú nulu na začiatku, aby sme vedeli vyjadriť aj počet pred všetkými zastávkami. Takže keby bolo na začiatku v električke povedzme \(5\) ľudí, tak po jednotlivých zastávkach by počty ľudí boli \(5, 7, 10, 3, 8, 4\).

Keď potom chceme vypočítať celkový počet nastupujúcich medzi nejakými dvoma zastávkami, tak je to rozdiel počtu, koľko ich bolo predtým, a koľko ich bolo potom. Tak napríklad medzi druhou a štvrtou zastávkou: po štvrtej zastávke bolo v električke o \(3\) ľudí viac ako na začiatku, a pred druhou zastávkou ich bolo o \(2\) viac ako na začiatku, teda musel nastúpiť jeden človek.

Všeobecne ak máme nejaký zoznam čísel a chceme vedieť rýchlo hľadať súčty na nejakých jeho úsekoch, hodí sa nám vypočítať si k nemu takéto “čiastkové súčty” – hovorí sa im prefixové súčty.

Na vyriešenie úlohy to nebolo treba, ale tu si ukážeme aj, ako by sa takéto niečo dalo naprogramovať. Dajme tomu, že v poli prefixove_sucty máme uložené prefixové súčty (začínajúce nulou). Potom súčet na nejakom úseku zistíme takto:

def sucet_na_useku(zaciatok, koniec):
    return prefixove_sucty[koniec + 1] - prefixove_sucty[zaciatok]

Podúloha B

Môžeme ich počítať postupne od začiatku a vždy využiť to, čo sme už vypočítali predtým. Ak nás zaujíma počet ľudí (rozdiel oproti začiatku) na \(k\)-tej zastávke, a už vieme, koľko ľudí tam nastúpilo, aj koľko ľudí v električke bolo predtým, tak nám stačí tieto čísla sčítať.

Začneme teda s tým, že na začiatku je rozdiel oproti počtu na začiatku \(0\). Teraz pre každú zastávku zoberieme rozdiel dovtedy (ktorý už máme vypočítaný), a pripočítame k nemu koľko ľudí nastúpilo na tej zastávke. Dostaneme tak rozdiel počtu na ďalšej zastávke oproti začiatku.

def spocitaj_prefixove_sucty(postupnost):
    sucty = [0]
    for pocet in postupnost:
        sucty.append(sucty[-1] + pocet)
    return sucty

Podúloha C

Súčet nastupujúcich na úseku je rozdiel počtu pred a po tom úseku, hľadáme teda taký začiatok a koniec v prefixových súčtoch, aby ich rozdiel bol čo najväčší. Dá sa vyskúšať všetky dvojice a zobrať tú s najväčším rozdielom, ale to by bolo zbytočne pomalé (ak máme \(n\) zastávok, museli by sme spraviť rádovo až \(n^2\) sčítaní).

Môžeme si ale uvedomiť, že keď si už nejakú zastávku vyberieme ako koniec, potom najväčší rozdiel bude mať s tou zastávkou predtým, na ktorej bolo najmenej ľudí. Stačí teda postupne prechádzať zastávky (tie budeme skúšať ako konce) a popritom si pamätať, aký najmenší počet ľudí bol doteraz na nejakej zasatávke (to bude ten začiatok). Takto sa určite dostaneme k takému začiatku a koncu, kde je ten rozdiel najväčší. Spravíme pri tom len rádovo \(n\) sčítaní.

def najvacsi_sucet_na_useku(postupnost):
    prefixove_sucty = spocitaj_prefixove_sucty(postupnost)
    najmenej_zatial = 0
    najvacsi_sucet = 0
    for pocet in prefixove_sucty:
        najvacsi_sucet = max(najvacsi_sucet, pocet - najmenej_zatial)
        najmenej_zatial = min(najmenej_zatial, pocet)
    return najvacsi_sucet

Podúloha D

Hľadáme, koľko je takých úsekov, že súčet nastupujúcich je tam \(0\), teda že predtým je tam rovnako veľa ľudí ako potom. Hľadáme teda počet rovnakých dvojíc v prefixových súčtoch.

Najprv si teda spočítame, koľkokrát sa v prefixových súčtoch každá hodnota vyskytuje. Potom pre každú jednu hodnotu nás zaujíma počet dvojíc prvkov s touto hodnotou. Ak je prvkov s touto hodnotou \(k\), potom počet dvojíc je \(k(k-1)/2\) (ak to nepoznáte, rozmyslite si, prečo to funguje). Stačí nám teda sčítať tieto počty dvojíc, na čo potrebujeme rádovo \(n\) sčítaní.

def pocet_nulovych_usekov(postupnost):
    prefixove_sucty = spocitaj_prefixove_sucty(postupnost)
    kolko_je_coho = defaultdict(int) # ako slovnik (dict), ale na zaciatku je ku kazdemu klucu priradena 0
    for pocet in prefixove_sucty:
        kolko_je_coho[pocet] += 1
    odpoved = 0
    for hodnota, pocet in kolko_je_coho.items():
        odpoved += pocet * (pocet - 1) // 2
    return odpoved