diff --git a/src/Init/Data/Array/Basic.lean b/src/Init/Data/Array/Basic.lean index 217451a403..5839247da3 100644 --- a/src/Init/Data/Array/Basic.lean +++ b/src/Init/Data/Array/Basic.lean @@ -21,6 +21,21 @@ def mkArray {α : Type u} (n : Nat) (v : α) : Array α := { data := List.replicate n v } +/-- +`ofFn f` with `f : Fin n → α` returns the list whose ith element is `f i`. +``` +ofFn f = #[f 0, f 1, ... , f(n - 1)] +``` -/ +def ofFn {n} (f : Fin n → α) : Array α := go 0 (mkEmpty n) where + /-- Auxiliary for `ofFn`. `ofFn.go f i acc = acc ++ #[f i, ..., f(n - 1)]` -/ + go (i : Nat) (acc : Array α) : Array α := + if h : i < n then go (i+1) (acc.push (f ⟨i, h⟩)) else acc +termination_by n - i + +/-- The array `#[0, 1, ..., n - 1]`. -/ +def range (n : Nat) : Array Nat := + n.fold (flip Array.push) (mkEmpty n) + @[simp] theorem size_mkArray (n : Nat) (v : α) : (mkArray n v).size = n := List.length_replicate .. @@ -413,6 +428,10 @@ def map {α : Type u} {β : Type v} (f : α → β) (as : Array α) : Array β : def mapIdx {α : Type u} {β : Type v} (as : Array α) (f : Fin as.size → α → β) : Array β := Id.run <| as.mapIdxM f +/-- Turns `#[a, b]` into `#[(a, 0), (b, 1)]`. -/ +def zipWithIndex (arr : Array α) : Array (α × Nat) := + arr.mapIdx fun i a => (a, i) + @[inline] def find? {α : Type} (as : Array α) (p : α → Bool) : Option α := Id.run <| as.findM? p @@ -487,6 +506,11 @@ def elem [BEq α] (a : α) (as : Array α) : Bool := def toList (as : Array α) : List α := as.foldr List.cons [] +/-- Prepends an `Array α` onto the front of a list. Equivalent to `as.toList ++ l`. -/ +@[inline] +def toListAppend (as : Array α) (l : List α) : List α := + as.foldr List.cons l + instance {α : Type u} [Repr α] : Repr (Array α) where reprPrec a _ := let _ : Std.ToFormat α := ⟨repr⟩ @@ -516,6 +540,13 @@ def concatMapM [Monad m] (f : α → m (Array β)) (as : Array α) : m (Array β def concatMap (f : α → Array β) (as : Array α) : Array β := as.foldl (init := empty) fun bs a => bs ++ f a +/-- Joins array of array into a single array. + +`flatten #[#[a₁, a₂, ⋯], #[b₁, b₂, ⋯], ⋯]` = `#[a₁, a₂, ⋯, b₁, b₂, ⋯]` +-/ +def flatten (as : Array (Array α)) : Array α := + as.foldl (init := empty) fun r a => r ++ a + end Array export Array (mkArray)