Największa suma mniejsza niż K

Dzisiaj zajmiemy się następującym problemem rekrutacyjnym:

Mamy daną tablicę A[1..N] oraz wartość K. Chcemy znaleźć spójną podtablicę A[i..j] o jak największej sumie, ale mniejszej od K.

Zacznijmy od pytania do rekrutera:

Kandydat: Czy tablica może zawierać elementy ujemne?
Rekruter: Tak.

To pytanie wynika z tego, że istnieje algorytm rozwiązujący problem w czasie liniowym dla tablic z wartościami nieujemnymi. Możecie się zastanowić nad nim w wolnym czasie.

Kandydat: Czy należy rozważyć pustą podtablicę?
Rekruter: Tak.
Kandydat: Czy możemy założyć, że K jest dodatnie?
Rekruter: Możemy.

Jeżeli odpowiedź na ostatnie pytanie byłaby negatywna – musielibyśmy się zapytać jaką wartość należy zwrócić jeśli wszystkie elementy tablicy są większe od K.

Rozwiązanie w czasie O(n³) otrzymujemy sprawdzając wszystkie możliwe wartości i oraz j:

int maksimum = 0;

for (int i = 0; i < n; i++)
    for (int j = i; j < n; j++)
    {
        int suma = 0;
        
        for (int l = i; l <= j; l++)
            suma += A[l];
        
        if (suma < K and suma > maksimum)
            maksimum = suma;
    }

Możemy zadanie w prosty sposób usprawnić, nie licząc wielokrotnie tych samych wartości (więcej o tym w notce o złożoności obliczeniowej):

int maksimum = 0;

for (int i = 0; i < n; i++)
{
    int suma = 0;
    
    for (int j = i; j < n; j++)
    {
        suma += A[j];
        
        if (suma < K and suma > maksimum)
            maksimum = suma;
    }
}

Otrzymujemy w ten sposób algorytm o złożoności O(n²).

Aby dalej ulepszyć nasze rozwiązanie, musimy przypomnieć sobie ideę sum prefiksowych. Zdefiniujmy sobie tablicę S[0..n] taką, że S[j] = A[1] + A[2] + ... + A[j] (dodatkowo niech S[0] = 0). Teraz, dla każdego j szukamy takiego i, żeby A[i] + A[i+1] + ... + A[j] było jak największą sumą, mniejszą od K. Ponieważ A[i] + A[i+1] + ... + A[j] = S[j] - S[i-1] zatem szukamy najmniejszej wartości S[i-1] większej od S[j] - K. Aby znaleźć takie elementy, możemy użyć BST, czyli seta z C++:

set<int> S;
int suma = 0;
int maksimum = 0;

S.insert(0);

for (int j = 0; j < n; j++)
{
    suma += A[j];
    S.insert(suma);
    
    int wynik = suma - S.lower_bound(suma - K + 1);
    
    if (wynik > maksimum)
        maksimum = wynik;
}

Algorytm działa w czasie O(n lg n) i używa O(n) dodatkowej pamięci.