aoc-swift/Sources/AOC2024/Day09.swift at develop · leohidalgo/aoc-swift · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import AOCCore
import Foundation

struct Day09: Day {
    let title = "Disk Fragmenter"
    var rawInput: String?

    func part1() throws -> Int {
        let diskMap = input().characters.compactMap(Int.init)
        var disk = makeDiskPart1(diskMap)

        var left = 0
        var right = disk.count - 1
        while left < right {
            while left < right, !disk[left].isFree { left += 1 }
            while left < right, disk[right].isFree { right -= 1 }

            disk.swapAt(left, right)
        }

        return disk.checksum()
    }

    func part2() throws -> Int {
        let diskMap = input().characters.compactMap(Int.init)
        var (files, free) = makeDiskPart2(diskMap)

        for id in (0..<files.count).reversed() {
            guard
                let (fileIndex, fileSize) = files[id],
                let i = free.firstIndex(where: { $0.size >= fileSize })
            else { continue }

            if fileIndex > free[i].index {
                files[id]?.index = free[i].index
                free[i].index += fileSize
                free[i].size -= fileSize
            }
        }

        return files.checksum()
    }

    private func makeDiskPart1(_ diskMap: [Int]) -> [Disk] {
        diskMap
            .enumerated()
            .reduce(into: (disk: [], id: 0)) { result, item in
                if item.offset.isEven {
                    result.disk.append(contentsOf: Array(repeating: Disk.file(id: result.id), count: item.element))
                    result.id += 1
                } else {
                    result.disk.append(contentsOf: Array(repeating: Disk.free, count: item.element))
                }
            }
            .disk as [Disk]
    }

    private func makeDiskPart2(_ diskMap: [Int]) -> (files: [Int: Block], free: [Block]) {
        let (files, free, _, _) = diskMap
            .enumerated()
            .reduce(into: (files: [Int: Block](), free: [Block](), id: 0, index: 0)) { result, item in
                if item.offset.isEven {
                    result.files[result.id] = (result.index, item.element)
                    result.id += 1
                } else {
                    result.free.append((result.index, item.element))
                }

                result.index += item.element
            }

        return (files, free)
    }
}

private typealias Block = (index: Int, size: Int)

private enum Disk: CustomDebugStringConvertible {
    case file(id: Int)
    case free

    var fileId: Int? {
        if case let .file(id) = self {
            id
        } else { nil }
    }

    var isFree: Bool {
        if case .free = self {
            true
        } else { false }
    }

    var debugDescription: String {
        switch self {
        case let .file(id): "\(id)"
        case .free: "."
        }
    }
}

private extension Dictionary where Key == Int, Value == Block {
    func checksum() -> Int {
        self
            .map {
                let startIndex = $0.value.index
                let endIndex = $0.value.index + $0.value.size

                return (startIndex..<endIndex).sum * $0.key
            }
            .sum
    }
}

private extension Collection where Element == Disk {
    func checksum() -> Int {
        self
            .compactMap(\.fileId)
            .enumerated()
            .map(*)
            .sum
    }
}