advent-of-code-2024/21/main.rb

#!/usr/bin/env ruby

require "debug"

input = (ARGV.first.nil? ? DATA : ARGF)
  .readlines(chomp: true)
  .map(&:chars)

@arrow_dirs = {
  "<" => [0, -1],
  "^" => [-1, 0],
  ">" => [0, 1],
  "v" => [1, 0],
}

@dir_arrows = @arrow_dirs.invert

numeric_movements = {
  "0" => ["^", ">"],
  "A" => ["<", "^"],
  "1" => ["^", ">"],
  "2" => ["<", "^", ">", "v"],
  "3" => ["<", "^", "v"],
  "4" => ["^", ">", "v"],
  "5" => ["<", "^", ">", "v"],
  "6" => ["<", "^", "v"],
  "7" => [">", "v"],
  "8" => ["<", ">", "v"],
  "9" => ["<", "v"],
}

directional_movements = {
  "^" => [">", "v"],
  "A" => ["v", "<"],
  "<" => [">"],
  "v" => ["^", ">", "<"],
  ">" => ["^", "<"],
}

numeric_keypad = [
  ["7", "8", "9"],
  ["4", "5", "6"],
  ["1", "2", "3"],
  [nil, "0", "A"],
]

directional_keypad = [
  [nil, "^", "A"],
  ["<", "v", ">"],
]

class PriorityQueue < Array
  alias_method :extract_min, :shift

  def <<(v)
    super

    sort_by!(&:last)
  end
end

def discover(char, graph)
  (0...graph.size).each do |x|
    (0...graph.first.size).each do |y|
      return [x, y] if char == graph.dig(x, y)
    end
  end
end

def dijkstra_paths(prev, target, current_path = [])
  return [current_path.reverse] if prev[target].nil?

  prev[target].flat_map do |parent|
    dijkstra_paths(prev, parent, current_path + [target])
  end
end

def dijkstra_with_lookup(graph, movements, start_char, target_char)
  start = discover(start_char, graph)
  target = discover(target_char, graph)

  dijkstra(graph, movements, start, target)
end

def dijkstra(graph, movements, start, target)
  queue = PriorityQueue.new
  dest = Hash.new
  prev = Hash.new

  queue << [start, 0]
  dest[start] = 0
  prev[start] = nil

  until queue.empty?
    node, moves = queue.extract_min

    movements[graph.dig(*node)]
      .map { @arrow_dirs[_1] }
      .map { |dx, dy| [node.first + dx, node.last + dy] }
      .each do |neighbor|
        new_moves = moves + 1
        if dest[neighbor].nil? || new_moves < dest[neighbor]
          dest[neighbor] = new_moves
          prev[neighbor] = [node]
          queue << [neighbor, new_moves]
        elsif new_moves == dest[neighbor]
          prev[neighbor] << node
        end
      end
  end

  paths = dijkstra_paths(prev, target)

  return dest, paths
end

@numeric_map = numeric_keypad
  .flatten.reject(&:nil?)
  .permutation(2).to_a
  .map { [[_1, _2], dijkstra_with_lookup(numeric_keypad, numeric_movements, _1, _2)] }
  .map do |(start, target), (dest, paths)|
    presses = paths.map do |path|
      ([discover(start, numeric_keypad)] + path)
        .each_cons(2).to_a
        .map { @dir_arrows[[_2.first - _1.first, _2.last - _1.last]] }
    end
    [[start, target], presses]
  end
  .to_h

@directional_map = directional_keypad
  .flatten.reject(&:nil?)
  .permutation(2).to_a
  .map { [[_1, _2], dijkstra_with_lookup(directional_keypad, directional_movements, _1, _2)] }
  .map do |(start, target), (dest, paths)|
    presses = paths.map do |path|
      ([discover(start, directional_keypad)] + path)
        .each_cons(2).to_a
        .map { @dir_arrows[[_2.first - _1.first, _2.last - _1.last]] }
    end
    [[start, target], presses]
  end
  .to_h

def minimum(chars, robots = 2, depth = 0, cache = {})
  return cache[[chars, depth]] if cache.key?([chars, depth])

  current, length, chars = "A", 0, Array(chars)
  return cache[[chars, depth]] = chars.map do |char|
    moves = (depth == 0 ? @numeric_map : @directional_map)
      .fetch([current, char], [[]])
      .map { _1 + ["A"] }
    current = char

    if depth == robots
      moves.map(&:size).min
    else
      moves.map { minimum(_1, robots, depth + 1, cache) }.min
    end
  end.sum
end

@numeric_map = @numeric_map
@directional_map = @directional_map

p input.map { minimum(_1) * _1.join.to_i }.sum
p input.map { minimum(_1, 25) * _1.join.to_i }.sum

__END__
029A
980A
179A
456A
379A