advent-of-code-2022/day-15/array-helpers.scm

(library (array-helpers)
  (export array-len-in-dim
          arrays->hash-table
          array-display
          array-map

          array-next-index
          array-cell-ref-vec
          array-index-of
          array-indices-of)
  (import
   (except (rnrs base)
           let-values
           map
           error
           vector-map)
   (only (guile)
         lambda* λ
         current-output-port
         ;; arrays
         array-shape
         array-ref
         array-cell-ref
         array-map!
         array-rank
         ;; display
         display
         simple-format
         current-output-port)
   ;; lists
   (srfi srfi-1)
   ;; vectors
   (srfi srfi-43)
   ;; hash tables
   (srfi srfi-69)
   (ice-9 arrays))


  (define array-len-in-dim
    (λ (arr dim)
      (let* ([shape (array-shape arr)]
             [dim-min-max (list-ref shape dim)])
        (+ (- (second dim-min-max)
              (first dim-min-max))
           1))))


  (define arrays->hash-table
    (lambda* (keys-arr vals-arrs #:optional (equal-func equal?))
      (let ([rows (array-len-in-dim keys-arr 0)]
            [cols (array-len-in-dim keys-arr 1)]
            [table (make-hash-table equal-func)])
        (let iter-rows ([row-ind 0])
          (let iter-cols ([col-ind 0])
            (cond
             [(< row-ind rows)
              (cond
               [(< col-ind cols)
                (hash-table-set! table
                                 (array-ref keys-arr row-ind col-ind)
                                 (array-ref vals-arrs row-ind col-ind))
                (iter-cols (+ col-ind 1))]
               [else (iter-rows (+ row-ind 1))])]
             [else table]))))))


  (define array-display
    (lambda* (landscape
              #:optional (port (current-output-port))
              #:key (formatter (λ (elem) elem)))
      (let ([rows (array-len-in-dim landscape 0)]
            [cols (array-len-in-dim landscape 1)])
        (let iter-rows ([row-ind 0])
          (let iter-cols ([col-ind 0])
            (cond
             [(>= row-ind rows) 'done]
             [(>= col-ind cols)
              (display "\n" (current-output-port))
              (iter-rows (+ row-ind 1))]
             [else
              (display (formatter (array-cell-ref landscape row-ind col-ind)) port)
              (iter-cols (+ col-ind 1))]))))))


  (define array-map
    (λ (proc src-arr)
      (define target-arr (array-copy src-arr))
      (array-map! target-arr proc src-arr)
      target-arr))


  (define array-next-index
    (λ (shape-vec indices-vec max-dim)
      "Increment one of indices in INDICES-VEC for which the
following conditions are true:

1. There is no index in INDICES-VEC that sits at a later
position (higher index) than MAX-DIM of INDICES-VEC and is
not yet at its maximum. The maximum is specified by the
SHAPE-VEC. Each index in INDICES-VEC has a corresponding
minimum and maximum in SHAPE-VEC at the same position.

After incrementing the index in INDICES-VEC, all later
indices (at a higher index of INDICES-VEC) are set to their
corresponding minimum to get a correct indices vector."
      (cond
       [(>= max-dim 0)
        (let ([index (vector-ref indices-vec max-dim)]
              [max-for-index (second (vector-ref shape-vec max-dim))])
          (cond
           [(< index max-for-index)
            ;; Copy the vector to not mutate argument.
            (let ([updated-indices-vec (vector-copy indices-vec)])
              ;; Increase index at position.
              (vector-set! updated-indices-vec
                           max-dim
                           (+ (vector-ref updated-indices-vec max-dim) 1))
              ;; Set later indices to their corresponding minimum.
              (let ([indices-vec-len (vector-length indices-vec)])
                (let iter ([dim° (+ max-dim 1)])
                  (cond
                   [(< dim° indices-vec-len)
                    (let ([minimum-ind-val (first (vector-ref shape-vec dim°))])
                      (vector-set! updated-indices-vec dim° minimum-ind-val)
                      (iter (+ dim° 1)))]
                   [else updated-indices-vec]))))]
           [else
            ;; Increment next higher dimension.
            (array-next-index shape-vec
                              indices-vec
                              (- max-dim 1))]))]
       [else #f])))


  (define array-cell-ref-vec
    (λ (arr indices-vec)
      "array-cell-ref takes a vector of indices INDICES-VEC, whose
length depends on how many dimensions the array ARR
has (what its rank is)."
      (let ([vec-len (vector-length indices-vec)])
        (let iter ([index-into-indices-vec° 0] [result arr])
          (cond
           [(< index-into-indices-vec° vec-len)
            (let ([cell-index (vector-ref indices-vec index-into-indices-vec°)])
              (iter (+ index-into-indices-vec° 1)
                    (array-cell-ref result cell-index)))]
           [else result])))))


  (define array-index-of
    (lambda* (arr pred #:optional (start-indices #f))
      "Return the index of the first element in ARR which
satisfies the predicate PRED."
      (let* ([shape (array-shape arr)]
             [shape-vec (list->vector shape)]
             [rank (array-rank arr)]
             [initial-indices
              (if start-indices
                  start-indices
                  (vector-map (λ (_i elem) (car elem)) shape-vec))])
        (let iter ([indices° initial-indices])
          (cond
           [indices°
            (cond
             ;; If the array element satisfies the
             ;; predicate, return the indices of the
             ;; element.
             [(pred (array-cell-ref-vec arr indices°))
              indices°]
             [else
              ;; Potential optimization: Make better use of
              ;; the rank argument, so that array-next-index
              ;; does not have to search unnecessarily for
              ;; the index to increment.
              (iter (array-next-index shape-vec indices° (- rank 1)))])]
           ;; No index found at which the predicate would be
           ;; satisfied.
           [else #f])))))


  (define array-indices-of
    (λ (arr pred)
      "Return a list of vectors of indices into the given array ARR which
satisfy the given predicate PRED, for an array of arbitrary rank."
      (let* ([shape (array-shape arr)]
             [shape-vec (list->vector shape)]
             [rank (array-rank arr)]
             [initial-indices (vector-map (λ (_i elem) (car elem)) shape-vec)])
        (let iter ([indices° initial-indices])
          (cond
           [indices°
            (cond
             [(pred (array-cell-ref-vec arr indices°))
              (cons indices°
                    (iter (array-next-index shape-vec indices° (- rank 1))))]
             [else (iter (array-next-index shape-vec indices° (- rank 1)))])]
           [else '()]))))))