4. Eeh?! Traktor na diaľnici?! vzorové riešenie

5
1
6 8 2 7 2

80

5
2
6 8 2 7 2

30

Naivné riešenie

Riešime úlohu, ako “rozrezať” súvislú skupinu úsekov \(u_1, u_2, \dots, u_N\) na \(K+1\) skupín tak aby celková kvalita bola čo najmenšia. Asi najednoduchší prístup je vyskúšať všetky možnosti rozdelenia a pre každé rozdelenie zrátať celkovú kvalitu. Všetky možnosti vyskúšame napríklad tak, že postupne budeme vyberať všetky kombinácie \(K\) prvkov spomedzi indexov \(1\) až \(N\). Pre jednu voľbu kombinácií \(c_1 < c_2 < \dots < c_K\) potom rozdelíme vstup do skupín \(\left( u_1, \dots , u_{c_1} \right), \left( u_{c_1+1}, \dots , u_{c_2} \right), \dots , \left( u_{c_K+1}, \dots , u_{N} \right)\), každej skupine zrátame jej skupinovú kvalitu a potom skupinové kvality sčítame do celkovej kvality.

Takéto riešenie je celkom pomalé - pôvodnú cestu rozdeľujeme celkovo \(N \choose K\) krát a v každom rozdelení rátame celkovú kvalitu, ktorú rátame ako súčet skupinových kvalít. Skupinovú kvalitu jednej skupiny vieme zrátať ako súčet súčinov všetkých dvojíc úsekov v skupine, čo trvá \(O(m^2)\) kde \(m\) je počet úsekov v skupine. Na zrátanie celkovej kvality teda minieme určite menej ako \(O(N^2)\) času, čím je celková časová zložitosť najviac \(O\left({N\choose K} \cdot N^2 \right)\).

S pamäťovou zložitosťou sme na tom lepšie, lebo ak máme dostatočne šikovnú implementáciu, tak si nám stačí pamätať iba pôvodné pole a indexy, ktoré vyberáme, čím je pamäťová zložitosť najviac \(O(N+K)\).

Odhad časovej zložitosti by sa dal vyrátať presnejšie, ale v tomto prípade na tom tak veľmi nezáleží. Vieme že časová zložitosť bude aspoň taká veľká ako kombinačné číslo, ktoré s rastúcim \(N, K\) rastie rýchlejšie ako hociaký polynóm, teda rozdiel v polynomiálnej konštante nemá taký veľký dopad. Napríklad \({2N \choose N} > 2^{N}\), pričom exponenciály rastú rýchlejšie ako polynómy.

Ak ste ešte nepočuli o kombinačných číslach, tak je dobrý čas sa o nich niečo naučiť napríklad na wikipédií

Dynamické riešenie

Skupinovú kvalitu jednej skupiny budeme rátať rovnako ako predtým, budeme avšak lepšie uvažovať nad tým, ako rozdeľovať úseky. Nech \(DP[s][k]\) odpovedá na nasledovnú otázku: Akú najmenšiu celkovú kvalitu viem dosiahnúť na úsekoch \(u_1, \dots , u_s\) ak sa Jožo povenuje presne \(k\) úsekom? Potom odpoveď ktorú hľadáme je rovná \(DP[N][K]\). Samozrejme nemá zmysel definovať \(DP[s][k]\) pre \(s < k\), keďže ani Jožo sa nevie povenovať viacerím úsekom ako je samotný počet úsekov. Zároveň nutne bude platiť, že \(DP[s][s-1] = DP[s][s]\), keďže už samotné \(DP[s][s-1]\) má celkovú kvalitu rovnú nule (každý úsek môže byť v samostatnej skupine). Stačí teda vyrátať hodnoty \(DP[s][k]\) pre \(s > k\) a budeme rovno vedieť aj hodnoty \(DP[s][s]\).

Pri rátaní \(DP[s][k]\) sa môžme riadiť veľmi užitočným pozorovaním - nikdy nemá zmysel venovať sa úseku \(u_s\), keďže namiesto neho sa môžem venovať inému úseku a znížiť tak hodnotu \(DP[s][k]\). Dôkaz je jednoduchý - keďže \(k < s\), tak existuje úsek, ktorému sa Jožo nepovenuje. Ak sa venuje úseku \(u_s\), tak si nijako nepomôže, teda je rozhodne lepšie venovať sa ľubovoľnému inému úseku (z predošlej vety je jasné, že si nájde iný úsek, ktorému sa môže venovať).

Na začiatku vieme zrátať všetky hodnoty tvaru \(DP[s][0]\) pre \(s \in \lbrace 1, \dots , N \rbrace\) jednoducho ako skupinové kvality úsekov \((u_1), (u_1, u_2),\) a tak ďalej až \((u_1, \dots, u_N)\). Ak už vieme hodnoty \(DP[i][k-1]\) pre \(i \in \lbrace 1, \dots , s-1\rbrace\), tak vieme zrátať \(DP[s][k]\) ako \(\min (DP[k][k-1] + C(k+1, s), DP[k+1][k-1] + C(k+2, s), \dots , DP[s-1][k-1] + C(s, s))\). Kde \(C(i, j)\) je skupinová kvalita úsekov \(u_i, \dots, u_j\).

Za týmto vzorcom sa schováva pomerne jednuduchá myšienka: V rozdelení do skupín pre každé \(DP[s][k]\) určite existuje taký úsek \(u_i\), že napravo od neho sa Jožo nevenoval žiadnemu ďalšiemu úseku. Môžme teda skúsiť všekty pozície takéhoto úseku - tj. že by to boli úseky \(u_{s-1}, u_{s-2}, \dots , u_k\). Úseky s menším indexom nemá zmysel skúšať, lebo pre nimi nie je \(k-1\) úsekov, ktorým sa môže Jožo venovať.

Takýto postup nám zaručí že v \(DP[N][K]\) budeme mať správnu odpoveď. Pamäťovú zložitosť budeme mať tento o čosi horšiu - potrebujeme \(O(N\cdot K)\) pamäte na to aby sme si zapamätali celú tabuľku \(DP\).

Ak budeme DP rátať za radom, to znamená že pred rátaním \(DP[s][k]\) vždy budeme mať zrátané všetky \(DP[a][b]\) pre \(a\leq s, b \leq k\), tak vždy budeme vedieť priamo zrátať \(DP[s][k]\) v čase \(O((s-k)^3) < O(s^3)\), lebo musíme vyskúšať \(s-k\) pozícií a pre každú pozíciu zrátať skupinovú cenu skupiny ktorá je veľká najviac \(s-k\). Celkovo teda máme riešenie v čas \(O(N^4\cdot K)\) na výpočet celej tabľky \(DP\) keďže \(s \leq N, k \leq K\).

Vzorová dynamika

Jediné čo zmeníme je výpočet skupinovej ceny úseku. Nech \(P[s]\) je súčet \(u_1 + \dots + u_s\) a \(Q[s]\) je súčet \(u_1^2 + \dots + u_s^2\). Potom vieme zistiť súčet alebo súčet štvorcov ľubovoľného úseku \(u_i, \dots, u_j\) ako \(P[j]-P[i-1]\) či \(Q[j]-Q[i-1]\). Skupinovú cenu úseku \(u_i, u_{i+1}, \dots, u_j\) potom zistíme ako \((P[j]-P[i-1])^2 - (Q[j]-Q[i-1])\) a to celé vydelíme \(2\). Funguje to podobne ako keď násobíme zátvorky - na zistenie súčtu \(ab+ac+bc\) môžme zrátať \((a+b+c)^2 = a^2+b^2+c^2+2(ab+ac+bc)\). Stačí teda odčítať štvorce a vydeliť dvomi.

Takýto postup má tú výhodu, že sčitať a násobiť čísla do určitej veľkosti zaberá konštantne veľa času, teda zlepšenie prišlo v tom, že na zrátanie \(DP[s][k]\) nám stačá konštantne málo základných aritmetických operácií, teda \(DP[s][k]\) vieme zrátať v čase \(O(s-k) = O(s)\).

Polia \(P, Q\) vieme spočítať v čase \(O(N)\), lebo z \(P[i], Q[i]\) vieme \(P[i+1], Q[i+1]\) spočítať ako \(P[i+1] = P[i]+u_{i+1}\), \(Q[i+1] = Q[i]+u_{i+1}^2\).

Celková časová zložitosť tak bude \(O(N^2\cdot K)\) (prejdeme \(O(NK)\) hodnôt \(Dp\), každú spracujeme v \(O(N)\)). Pamäťová zložitosť bude \(O(NK)\).

#!/usr/bin/env python3

#input
N = int(input())
K = int(input())
Useky = [int(x) for x in input().split()]

#spočítame prefixové polia P, Q:
pref_sum = [0]
pref_sqr = [0]
for i in range(0, len(Useky)):
    pref_sum.append(pref_sum[-1]+Useky[i])
    pref_sqr.append(pref_sqr[-1]+Useky[i]*Useky[i])

def Cost(a, b):
    '''vracia skupinovú kvalitu skupiny u_a, ____ , u_b'''
    if a > b:
        return 0
    return ((pref_sum[b+1]-pref_sum[a])**2 - (pref_sqr[b+1]-pref_sqr[a])) // 2

# DP[s][k] zo vzoráku
DP = [[10**18 for _ in range(K+1)] for _ in range(N)]
for i in range(len(Useky)):
    DP[i][0] = Cost(0, i)

# treba si dať pozor na to, že vo vzoráku úseky číslujeme $u_1, ____ , u_N$,
# avšak polia sa číslujú od 0 - treba teda väčšie pole alebo zmenšiť príslušné indexy.
for k in range(1, K+1):
    for s in range(0, len(Useky)):
        for i in range(k-1, s):
            DP[s][k] = min(DP[s][k], DP[i][k-1] + Cost(i+1, s))

print(DP[N-1][K])