Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- def solve() -> None:
- disk_map: list[int] = list(map(int, get_string(as_list=True)))
- disk_map_with_dots: list[int] = []
- curr_id: int = 0
- heaps_free_space: list[list[int]] = [[] for _ in range(10)] # one heapq per size
- # in the heapq we store idx first free space
- for i in range(len(disk_map)):
- if i % 2 == 0:
- # it's a file
- disk_map_with_dots += [curr_id for _ in range(disk_map[i])]
- curr_id += 1
- else:
- # it's free space
- if disk_map[i] > 0:
- heapq.heappush(heaps_free_space[disk_map[i]], len(disk_map_with_dots))
- for _ in range(disk_map[i]):
- disk_map_with_dots.append(-1)
- i: int = len(disk_map_with_dots) - 1
- while i >= 0:
- if disk_map_with_dots[i] != -1:
- file_id: int = disk_map_with_dots[i]
- file_width: int = 0
- while i >= 0 and disk_map_with_dots[i] == file_id:
- i -= 1
- file_width += 1
- # swap with left most
- smallest_idx: int = 10**18
- best_width: int = 0
- for width in range(file_width, 10):
- if heaps_free_space[width]:
- if smallest_idx > heaps_free_space[width][0]:
- smallest_idx = heaps_free_space[width][0]
- best_width = width
- if smallest_idx != 10**18 and smallest_idx <= i:
- for j in range(file_width):
- disk_map_with_dots[smallest_idx + j] = file_id
- disk_map_with_dots[i + j + 1] = -1
- heapq.heappop(heaps_free_space[best_width])
- if best_width - file_width > 0:
- heapq.heappush(heaps_free_space[best_width - file_width], smallest_idx + file_width)
- else:
- i -= 1
- ans: int = 0
- for i in range(len(disk_map_with_dots)):
- if disk_map_with_dots[i] != -1:
- ans += disk_map_with_dots[i] * i
- print(ans)
- if __name__ == "__main__":
- import time
- start_time = time.time()
- solve()
- print("--- %s seconds ---" % (time.time() - start_time))
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement