Dzisiejsze zadanie podniosło nieco poziom trudności i sprawiło że musiałem dłużej nad nim pomyśleć.

Treść

Naszym zadaniem jest dokonać defragmentacji dysku. Na wejściu dostajemy listę cyfr, które oznaczają ile bloków dysku zajętych jest przez pliki i puste przestrzenie. Cyfry na parzystych pozycjach oznaczają ilość bloków zajętych przez pliki, zaś na nieparzystych oznaczają puste bloki.

1234501   # tak to wygląda na wejściu
0..111....22222.3  # tak to wygląda na dysku (. to puste miejsca a cyfry to idki plików)

Każdy plik ma swoje id i są one przydzielane inkrementując id o jeden przy każdym kolejnym zajętym klastrze bloków z wejścia.

Naszym zadaniem jest poprzestawiać bloki tak żeby nie było żadnej pustej przestrzeni pomiędzy plikami, zaczynamy od bloków położonych najbardziej na prawo i przenosimy je najbardziej na lewo jak się da, do powyższego przykładu rozwiązanie będzie takie:

0321112222......

Odpowiedzią jest checksuma całego dysku którą obliczamy mnożąc idki pliku poszczególnych bloków z ich pozycją w dysku i sumując wszystkie wartości, czyli:

checksum = 0*0 + 1*3 + 2*1 + 3*1 + ...

Rozwiązanie

W moim rozwiązaniu trzymam dwa wskaźniki, lewy na początku programu ustawiony jest na pierwszy blok, z kolei prawy na ostatni blok. Będziemy przenosić bloki tak długo aż lewy wskaźnik nie wyprzedzi prawego. Jeśli lewy wskaźnik będzie znajdował się na fragmencie z plikiem, doda ten fragment do checksumy, jeśli stanie na pustym fragmencie to wtedy będzie dodawał do checksumy fragment na który wskazuje prawy wskaźnik (w takiej ilości jaka się zmieści)

disk = list(map(int, input()))

left = 0
right = len(disk) - 1

full = True  # jeśli jesteśmy na bloku z plikiem to True, jeśli na pustym to False
offset = 0   # ile bloków zostało już wliczonych do checksumy
checksum = 0

while left <= right:
    n = disk[left]
    if full:
        multiplier = (n * (n - 1)) / 2 + offset * n  # używamy wzoru na liczby trójkątne

        checksum += multiplier * (left // 2)  # left // 2 daje nam id bloku
        offset += n
        left += 1
        full = False
    else:
        last_n = disk[right]
        to_move = min(disk[right], n)
        disk[left] -= to_move
        disk[right] -= to_move

        multiplier = (to_move * (to_move - 1)) / 2 + offset * to_move
        checksum += multiplier * (right // 2)  # right // 2 daje nam id bloku

        offset += to_move
        if disk[right] == 0:
            right -= 2  # prawy pomija wszystkie puste bloki
        if disk[left] == 0
            left += 1
            full = True

print(checksum)

Rozwiązanie może nie jest najczystsze ale jest bardzo szybkie i nie zajmuje żadnej dodatkowej pamięci poza wejściem programu.

Wzór na liczbę trójkątną

Jeśli wiemy że musimy obliczyć fragment checksumy dla pliku który zaczyna się w indeksie o i ma długość n, to możemy to otrzymamy taką wartość:

(o + (o + 1) + (o + 2) + ... + (o + n - 1)) * id

Możemy całkiem sprytnie wyciągnąć sobie wzór na ten ciąg, dajmy przykład gdzie: o=3, n=3, wtedy mamy coś takiego:

  |
 ||  # trójkąt   (0 + 1 + 2 = (n * (n - 1)) // 2 )
|||  # prostokąt (3 * 3     = o * n)
|||
|||  

czyli końcowy wzór: (n * (n - 1)) // 2 + on.

Część druga

Tutaj sprawa się lekko komplikuje, okazuje się że możemy przenieść dany plik tylko jeśli wszystkie jego bloki zmieszczą się w pustej przestrzeni (nadal przneosimy najpierw prawe pliki jak najbardziej na lewo).

Dokonam dodatkowego preprocessingu, żeby łatwiej pisać dalszą część:

disk = list(map(int, input()))
offsets = list(accumulate([0] + disk))[:-1]
D = list(zip(disk, offsets))
full = D[::2]
empty = D[1::2]

offsets - zawiera indeksy na których zaczyna się dany fragment dysku. full - lista tupli (offset, ilość) która zawiera framgnety wypełnione plikami empty - podobnie jak full ale zawiera puste przestrzenie

Jeśli chcemy żeby nasze rozwiązanie było optymalne to musimy upewnić się że znajdywanie odpowiednich pustych fragmentów dla nowych plików będzie jak najszybsze, jedną obserwacją jest to że mamy tylko 9 możliwych wielkości pustych przestrzeni, więc możemy mapować indeksy listy do tupli z pustymi przestrzeniami o odpowiedniej wielkości, chcemy też żeby zawzsze rozważać przestrzenie znajdujące się jak najbardziej w lewo dlatego użyjemy kopca by móc zawsze odzyskać przestrzeń z najmniejszym offsetem. (przyda się to bo będziemy też dodawać nowe puste przestrzenie jeśli zapełnimy je tylko częściowo).

Tworzenie indeksu

index = []
for i in range(10):
    index.append([])

for n, eoffset in empty:
    heapq.heappush(index[n], eoffset)

Reszta kodu:

checksum = 0
for identifier in reversed(range(len(full))):  # iterujemy po blokach od tyłu
    moved = False
    n, offset = full[identifier]

    # szukamy najbardziej lewego pustego bloku o odpowiedniej wielkości
    leftmost = float("inf")
    chosen = None
    for en in range(n, len(index)):
        if index[en] and index[en][0] < leftmost and offset > index[en][0]:
            leftmost = index[en][0]
            chosen = en

    # jeśli nie znaleźliśmy pustego bloku to liczymy fragment checksumy bez przesunięcia
    if not chosen:
        checksum += ((n * (n - 1)) // 2 + offset * n) * identifier
        continue

    eoffset = heapq.heappop(index[chosen])
    checksum += ((n * (n - 1)) // 2 + eoffset * n) * identifier
    diff = chosen - n
    if diff:  # nie zapełniliśmy całej przestrzeni więc dodajemy nowy wpis do indeksu
        heapq.heappush(index[diff], eoffset + n)
    moved = True

print(checksum)

Po uruchomieniu programu dostajemy poprawną odpowiedź do drugiej części, działa to bardzo szybko, na moim laptopie wykonuje się w 0.029 sekundy.

Podsumowanie

Byłem w stanie znaleźć rozwiązania które się szybko wykonują, choć poświęciłem na to więcej czasu niż chciałbym przyznać. Domyślam się że bardziej bruteforcowe rozwiązanie mógłbym napisać w dużo krótszym czasie, ale na pewno byłaby mniejsza satysfakcja z wyniku.