Arrayify

247.3ms ± 4.6ms shakespeare

Main.lean

@@ -1,10 +1,14 @@
 def String.words (s : String) :=
-  s.split Char.isWhitespace |>.filter (¬·.isEmpty)
+  s.split Char.isWhitespace |>.filter (¬·.isEmpty) |>.toArray
 def String.lines (s : String) :=
-  s.splitOn "\n" |>.filter (¬·.isEmpty)
+  s.splitOn "\n" |>.filter (¬·.isEmpty) |>.toArray
 
-def List.unlines (xs : List String) : String
-  := xs.intersperse "\n" |> String.join
+-- TODO: remove trailing
+def Array.unlines (xs : Array String) : String := Id.run do
+  let mut out := ""
+  for line in xs do
+    out := out ++ line ++ "\n"
+  return out
 
 def ex := "The heat bloomed\tin December\nas the carnival  season\n        kicked into gear.\nNearly helpless with sun and glare, I avoided Rio's brilliant\nsidewalks\n  and glittering beaches,\npanting in dark   corners\nand waiting out the inverted southern summer."
 
@@ -12,56 +16,57 @@ def ex := "The heat bloomed\tin December\nas the carnival  season\n        kicke
 #eval ex.lines
 #eval IO.println ex.lines.unlines
 
-def sep : List a → List a → List a
-| a1::a1s, a2::a2s => a1::a2::sep a1s a2s
-| a1,       _      => a1
+def Array.replicate : Nat → α → Array α
+| 0,   _ => #[]
+| n+1, x => replicate n x |>.push x
 
-def justify : Nat → List String → String
-| _,    [word] => word
+theorem Array.replicate_size_eq_n : (Array.replicate n x).size = n := by
+  induction n <;> simp [Array.replicate, *]
+
+def String.replicate : Nat → Char → String
+| 0,   _ => ""
+| n+1, x => replicate n x |>.push x
+
+def sep (xs space : Array String)
+  : String := Id.run do
+  let mut out := ""
+  let mut i := 0
+  while h : i < xs.size do
+    out := out.append xs[i]
+    if let some space := space[i]?
+    then out := out.append space
+    i := i + 1
+  return out
+
+def justify : Nat → Array String → String
+| _,   #[word] => word
 | width, words =>
-    let wordsLength   := List.sum (String.length <$> words)
+    let wordsLength   := words.foldr (·.length + ·) 0
     let spacesNeeded  := width - wordsLength
-    let numSpaces     := words.length - 1
+    let numSpaces     := words.size - 1
     let smallSpaceLen := spacesNeeded / numSpaces
-    let smallSpace    := List.replicate smallSpaceLen ' '
+    let smallSpace    := String.replicate smallSpaceLen ' '
     let largeSpaces   := spacesNeeded % numSpaces
     let smallSpaces   := numSpaces - largeSpaces
-    let largeSpace    := ' '::smallSpace
-    let spaces := List.replicate largeSpaces largeSpace
-               ++ List.replicate smallSpaces smallSpace
-               |>.map List.asString
-    String.join $ sep words spaces
+    let largeSpace    := smallSpace.push ' '
+    let spaces := Array.replicate largeSpaces largeSpace
+               ++ Array.replicate smallSpaces smallSpace
+    sep words spaces
 
 #eval justify 1 "Hello, world, how is it going?".words
 #eval justify 1 "Hello, world, how is it going?".words |>.length
 
-def breakUp (width : Nat) (words : List String) : List (List String) :=
+def breakUp (width : Nat) (words : Array String) : Array (Array String) :=
   let (_, line, lines) := words.foldl
     (λ(left, line, lines) word ↦
       if word.length < left
-      then (left - (word.length + 1), line++[word], lines)
-      else (width - word.length, [word], lines++[line]))
-    (width, [], [])
-  lines++[line]
-
--- def breakUp (width : Nat) (words : List String)
---   : List (List String) :=
---   if width = 0 then [] else go width words [] []
---   where
---   go (left : Nat) (words line : List String) (lines : List $ List String) :=
---   match left, words, line, lines with
---   | _,    [],      l, ls => ls++[l]
---   | left, (w::ws), l, ls =>
---       if width < w.length
---       then go width ws [w] ls
---       else if left < w.length + 1
---       then go width (w::ws) [] (ls++[l])
---       else go (left - (w.length + 1)) ws (l++[w]) ls
---   termination_by (left, words)
-
+      then (left - (word.length + 1), line.push word, lines)
+      else (width - word.length, #[word], lines.push line))
+    (width, #[], #[])
+  lines.push line
 
 def solve (width : Nat) :=
-  List.unlines ∘ List.map (justify width) ∘ breakUp width ∘ String.words
+  Array.unlines ∘ Array.map (justify width) ∘ breakUp width ∘ String.words
 
 #eval IO.println $ solve 40 ex