!YS-v0

defn main(n=10):
  gprint:
    zig-zag: n

defn gprint(grid):
  f =: (" %$(grid.$.$.digits().#)d" * grid.#)
  each line grid:
    say: format(f line*)

defn zig-zag(n):
  take-from n:
    map \(_1 _2) cycle([reverse a]):
      partitions _ range(n:sqr):
        concat (1 .. n): (n.-- .. -1)

defn take-from(n colls):
  lazy-seq:
    when-let s seq(colls):
      first-n rest-n =: split-at(n s)
      cons map(first first-n):
        take-from n:
          concat _ rest-n:
            filter seq:
              map rest: first-n

defn partitions(sizes coll):
  lazy-seq:
    when-let n sizes.0:
      when-let s seq(coll):
        cons take(n coll):
          partitions next(sizes):
            drop n: coll