advent-of-code/2019/06/src/main.zig

const std = @import("std");

const Link = struct {
    a: [3]u8 = [3]u8{ 0, 0, 0 },
    b: [3]u8 = [3]u8{ 0, 0, 0 },
};

const Node = struct {
    name: [3]u8 = [3]u8{ 0, 0, 0 },
    depth: u32 = 0,
    parent: ?*Node,
};

const Links = std.ArrayList(Link);
const NodeList = std.ArrayList(*Node);
const Nodes = std.StringHashMap(*Node);

fn parse(alloc: *std.mem.Allocator, stream: *std.fs.File.InStream.Stream) !Links {
    var list: Links = Links.init(alloc);
    var buf: [8]u8 = undefined;
    var i: u32 = 0;

    while (try stream.readUntilDelimiterOrEof(&buf, '\n')) |line| {
        std.debug.assert(line.len == 7 and line[3] == ')');

        var out: Link = Link{};
        std.mem.copy(u8, &out.a, line[0..3]);
        std.mem.copy(u8, &out.b, line[4..7]);

        try list.append(out);
    }

    return list;
}

fn getOrCreateNode(alloc: *std.mem.Allocator, nodes: var, name: [3]u8) !*Node {
    if (nodes.get(name)) |kv| {
        return kv.value;
    } else {
        var node: *Node = try alloc.create(Node);
        node.parent = null;
        node.depth = 0;

        std.mem.copy(u8, &node.name, name);

        var kv = try nodes.put(name, node);
        return node;
    }
}

fn getParents(alloc: *std.mem.Allocator, node: *Node) !NodeList {
    var out = NodeList.init(alloc);
    var n: *Node = node;
    while (n.parent) |p| {
        try out.append(p);
        n = p;
    }

    return out;
}

// Assumes n1 is a parent of n2
fn distance(n1: *Node, n2: *Node) u32 {
    return n2.depth - n1.depth - 1;
}

pub fn main() anyerror!void {
    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();

    const alloc = &arena.allocator;
    const file = std.io.getStdIn();
    const stream = &file.inStream().stream;

    var links = try parse(alloc, stream);

    // name -> node
    var nodes = Nodes.init(alloc);

    // The input isn't ordered, so we need to pass over it several times

    // Pass 1: create all nodes
    for (links.toSlice()) |link| {
        var a = try getOrCreateNode(alloc, &nodes, link.a);
        var b = try getOrCreateNode(alloc, &nodes, link.b);
    }

    // Pass 2: link all nodes to their parents
    for (links.toSlice()) |link| {
        var a = try getOrCreateNode(alloc, &nodes, link.a);
        var b = try getOrCreateNode(alloc, &nodes, link.b);

        std.debug.assert(b.parent == null);
        b.parent = a;
    }

    // Pass 3: The tree is fully constructed, so we can count parents
    var count: u32 = 0;
    var iter = nodes.iterator();
    while (iter.next()) |kv| {
        var node = kv.value;

        // Count number of orbits
        var n: *Node = node;
        while (n.parent) |p| {
            node.depth += 1;
            n = p;
        }

        count += node.depth;
    }

    std.debug.warn("Part 1: {}\n", count);

    // We need to find the node that both you and santa have in common

    var you = (nodes.get("YOU") orelse unreachable).value;
    var san = (nodes.get("SAN") orelse unreachable).value;

    var yP = try getParents(alloc, you);
    var sP = try getParents(alloc, san);

    var inCommon = NodeList.init(alloc);

    for (yP.toSlice()) |p1| {
        for (sP.toSlice()) |p2| {
            if (p1 == p2) try inCommon.append(p1);
        }
    }

    std.debug.warn("Part 2: you have {} parents, santa has {}, {} in common\n", yP.count(), sP.count(), inCommon.count());

    for (inCommon.toSlice()) |p| {
        const d1 = distance(p, you);
        const d2 = distance(p, san);
        std.debug.warn("\t{} - ({}) {} ({}) - {}: {}\n", you.name, d1, p.name, d2, san.name, d1 + d2);
    }
}