94 lines
2.7 KiB
Text
94 lines
2.7 KiB
Text
/-
|
||
Copyright (c) 2023 Lean FRO, LLC. All rights reserved.
|
||
Released under Apache 2.0 license as described in the file LICENSE.
|
||
Authors: Joachim Breitner
|
||
-/
|
||
module
|
||
|
||
prelude
|
||
import Init.Data.Stream
|
||
public import Init.Data.Range.Polymorphic.Nat
|
||
public import Init.Data.Range.Polymorphic.Iterators
|
||
|
||
namespace Array
|
||
|
||
/-!
|
||
This module contains utility functions involving Arrays that are useful in a few places
|
||
of the lean code base, but too specialized to live in `Init.Data.Array`, which arguably
|
||
is part of the public, user-facing standard library.
|
||
-/
|
||
|
||
/--
|
||
Compares each element of an array with all later elements using `f`. For each comparison, `f`
|
||
determines whether to keep both of its arguments. If `f` returns `false` for an argument, that
|
||
argument is removed from the array and does not participate in subsequent comparisons. Those
|
||
elements that were not discarded are returned.
|
||
|
||
This can be used to remove elements from an array where a “better” element, in some partial order,
|
||
exists in the array.
|
||
|
||
Example:
|
||
```lean example
|
||
#eval #["a", "r", "red", "x", "r"].filterPairsM fun x y =>
|
||
pure (!(x.isPrefixOf y), true)
|
||
```
|
||
```output
|
||
#["a", "red", "x", "r"]
|
||
```
|
||
-/
|
||
public def filterPairsM {m} [Monad m] {α} (a : Array α) (f : α → α → m (Bool × Bool)) :
|
||
m (Array α) := do
|
||
let mut removed := Array.replicate a.size false
|
||
let mut numRemoved := 0
|
||
for h1 : i in *...a.size do for h2 : j in (i+1)...a.size do
|
||
unless removed[i]! || removed[j]! do
|
||
let xi := a[i]
|
||
let xj := a[j]
|
||
let (keepi, keepj) ← f xi xj
|
||
unless keepi do
|
||
numRemoved := numRemoved + 1
|
||
removed := removed.set! i true
|
||
unless keepj do
|
||
numRemoved := numRemoved + 1
|
||
removed := removed.set! j true
|
||
let mut a' := Array.mkEmpty numRemoved
|
||
for h : i in *...a.size do
|
||
unless removed[i]! do
|
||
a' := a'.push a[i]
|
||
return a'
|
||
|
||
/--
|
||
`maskArray mask xs` keeps those `x` where the corresponding entry in `mask` is `true`
|
||
-/
|
||
public def mask {α} (mask : Array Bool) (xs : Array α) : Array α := Id.run do
|
||
let mut ys := #[]
|
||
for b in mask, x in xs do
|
||
if b then ys := ys.push x
|
||
return ys
|
||
|
||
/--
|
||
Inverse of `Array.mask`:
|
||
```
|
||
Array.zipMasked mask (Array.mask (mask.map not) xs) (Array.mask mask xs) == xs
|
||
```
|
||
-/
|
||
public def zipMasked {α} (mask : Array Bool) (xs ys : Array α) : Array α := Id.run do
|
||
let mut i := 0
|
||
let mut j := 0
|
||
let mut zs := #[]
|
||
for b in mask do
|
||
if b then
|
||
if h : j < ys.size then
|
||
zs := zs.push ys[j]
|
||
j := j + 1
|
||
else
|
||
panic! "zipMaskedArray: not enough elements in ys"
|
||
else
|
||
if h : i < xs.size then
|
||
zs := zs.push xs[i]
|
||
i := i + 1
|
||
else
|
||
panic! "zipMaskedArray: not enough elements in xs"
|
||
return zs
|
||
|
||
end Array
|