1i [
1,$ {
	s/^/[/
	s/$/]/
}
2,$i,
$i ]
$d

Basic

import "lib/github.com/diku-dk/sorts/radix_sort"

type position = (u32, u32, u32)
def positionFromArray (p: [3]u32): position
  = (p[0], p[1], p[2])
def pair_iota (n: i64): [n](i64, i64)
  = map (\ j -> (n, j)) (iota n)
def gaussian_sum (n: i64) = n * (n + 1) / 2

def euclidean_distance (a: position) (b: position): f64
  = f64.sqrt 
    ( (f64.u32 a.0 - f64.u32 b.0) ** 2
    + (f64.u32 a.1 - f64.u32 b.1) ** 2
    + (f64.u32 a.2 - f64.u32 b.2) ** 2
    )

def distance_table [n] (positions: [n]position): [n][n]f64
  = let distance_function = \ i j -> euclidean_distance positions[i] positions[j] in
    tabulate_2d n n distance_function

def existsLength 'a 'b (f: a -> ?[l].[l]b) (x: a): i64
  = length (f x)

def concatMap [n] 'a 'b (f: a -> ?[l].[l]b) (placeholder: b) (xs: [n]a): *[]b
  = let totalLength = reduce (+) 0 <| map (\ x -> length (f x)) xs in
    ( loop (results, offset) = (replicate totalLength placeholder, 0)
      for x in xs
      do
        let bs = f x in
        let scatterIndices = indices bs |> map (+offset) in
        (scatter results scatterIndices bs, offset + length bs)
    ).0

def distance_array [n] (positions: [n]position): []((i64, i64), f64)
  = let table = distance_table positions in
    let triangle_indices = concatMap pair_iota (i64.highest, i64.highest) (iota n |> drop 1) in
    map (\ (i, j) -> ((i, j), table[i, j])) triangle_indices

def sort_distances (distances: []((i64, i64), f64)): []((i64, i64), f64)
  = radix_sort_float_by_key (.1) f64.num_bits f64.get_bit distances

type option 'a
  = #Empty
  | #Present a

def empty 'a : option a = #Empty

def overrideWith (old: u16) (new: u16) (x: option u16): option u16
  = match x
      case #Empty -> #Empty
      case #Present inner -> 
        if inner == old
        then #Present new
        else #Present inner

def orElse 'a (o: option a) (d: a): a
  = match o
      case #Empty -> d
      case #Present x -> x

def is_present 'a (o: option a): bool
  = match o
      case #Empty -> false
      case #Present _ -> true

def connect (circuits: *[](option u16)) (newCircuitId: u16) (connection: (i64, i64)): (u16, *[](option u16))
  = let circuitA = circuits[connection.0] in
    let circuitB = circuits[connection.1] in
    match (circuitA, circuitB)
      case (#Empty, #Empty) -> 
        ( newCircuitId + 1
        , scatter circuits [connection.0, connection.1] (rep (#Present newCircuitId))
        )
      case (#Present a, #Empty) -> 
        ( newCircuitId
        , scatter circuits [connection.1] [#Present a]
        )
      case (#Empty, #Present b) -> 
        ( newCircuitId
        , scatter circuits [connection.0] [#Present b]
        )
      case (#Present a, #Present b) ->
        ( newCircuitId
        , map (b `overrideWith` a) circuits
        )

def countCircuit (counts: *[]u64) (o: option u16): *[]u64 
  = match o
    case #Empty -> counts 
    case #Present i -> scatter counts [i64.u16 i] [counts[i64.u16 i] + 1]

def countCircuits (c: u16) (circuits: [](option u16)): *[i64.u16 c]u64
  = let circuitCounts = replicate (i64.u16 c) 0 in
    loop counts = circuitCounts
    for circuit in circuits
    do countCircuit counts circuit

def exampleConnectionCount = 10i64
def inputConnectionCount = 1000i64

def part1 (positions: i64) (connectionCount: i64) (distances: []((i64, i64), f64))
  = let connections = take connectionCount distances |> map (.0) in
    let circuitMap: *[positions](option u16) = replicate positions empty in
    ( loop (circuitCount, circuits) = (0, circuitMap)
      for connection in connections
      do
        connect circuits circuitCount connection
    ) |> uncurry countCircuits 
      |> radix_sort u64.num_bits u64.get_bit
      |> reverse
      |> take 3
      |> foldl (*) 1

def part2 (positionCount: i64) (distances: []((i64, i64), f64)) (positions: []position)
  = let circuitMap: *[positionCount](option u16) = replicate positionCount empty in
    ( loop (circuitCount, connectionIndex, circuits) = (0, 0, circuitMap)
      while not
        ( and (map is_present circuits)
        && and (map (== circuits[0]) circuits)
        )
      do
        let connection = distances[connectionIndex].0 in
        let (newCircuitId, circuits') = connect circuits circuitCount connection in
        (newCircuitId, connectionIndex+1, circuits')
    ).1
    |> \ i -> distances[i-1].0
    |> \ (a, b) -> positions[a].0 * positions[b].0

def main [n] (position_array: [n][3]u32)
  = let positions = map positionFromArray position_array in
    let unsorted_distances = distance_array positions in
    let sorted_distances = sort_distances unsorted_distances in
    ( part1 n inputConnectionCount sorted_distances
    , part2 n sorted_distances positions
    )