advent_of_code_2022_day_7.nim

import std/sequtils, std/strutils, std/sugar, std/strformat

type 
    NodeKind = enum File, Dir

    Node = object
        kind: NodeKind
        name: string
        size: int
        children: seq[Node]

proc buildTree(input: seq[string]): Node =
    var stack: seq[ptr Node]
    var root = Node(kind: Dir, name: "/")

    for line in input:
        if line == "$ cd /":
            stack = @[addr root]
        elif line == "$ cd ..":
            stack.delete(stack.len-1)
        elif line.startsWith("$ cd "):
            let dirName = line["$ cd ".len..<line.len]
            var currentRoot = stack[^1]

            for idx in 0..<currentRoot.children.len:
                var node = currentRoot.children[idx]

                if node.name == dirName and node.kind == Dir:
                    stack.add(addr currentRoot.children[idx])
                    break
        elif not line.startsWith("$ "):
            let parts = line.split(" ")
            let name = parts[1]

            if parts[0] == "dir":
                stack[stack.len-1].children.add(
                    Node(kind: Dir, name: name)
                )
            else:
                let size = parseInt(parts[0])

                for idx in 0..<stack.len:
                    stack[idx].size += size
                
                stack[stack.len-1].children.add(
                    Node(kind: File, name: name, size: size)
                )

    root

proc traverseTree[T](node: Node, initial: T, procFun: (Node, T) -> T): T =
    result = initial
    var stack = @[node]
    var currentNode: Node

    while stack.len > 0:
        currentNode = stack.pop()
        result = procFun(currentNode, result)
        for child in currentNode.children:
            stack.add(child)

proc sumDirsUpTo100k(node: Node, acc: int): int =
    if node.kind == Dir and node.size <= 100000:
        result = acc + node.size
    else:
        result = acc

proc smallestDirSizeAbove(sizeToFreeUp: int): (Node, int) -> int =
    result = proc(node: Node, size: int): int =
        if node.kind == Dir and node.size >= sizeToFreeUp and node.size < size:
            node.size
        else:
            size

let input = "day-7/input.txt".lines.toSeq

let tree = buildTree(input)

let spaceToFreeUp = 30000000 - (70000000 - tree.size)

echo fmt"Sum of size of directories at or below 100k: {traverseTree(tree, 0, sumDirsUpTo100k)}"
echo fmt"Smallest dir size to free up 30M of space: {traverseTree(tree, high(int), smallestDirSizeAbove(spaceToFreeUp))}"