57 lines
1.5 KiB
Haskell
57 lines
1.5 KiB
Haskell
|
#!/usr/bin/env nix-script
|
||
|
|
#!>haskell
|
||
|
|
#! haskell | vector
|
||
|
|
|
||
|
|
{- Solve a linear system by Gauss elimination. -}
|
||
|
|
|
||
|
|
import Data.Vector (Vector, cons, snoc, empty)
|
||
|
|
import qualified Data.Vector as V
|
||
|
|
|
||
|
|
type Matrix a = Vector (Vector a)
|
||
|
|
|
||
|
|
decons :: Vector a -> (a, Vector a)
|
||
|
|
decons x = (V.head x, V.tail x)
|
||
|
|
|
||
|
|
asList :: ([a] -> [b]) -> (Vector a -> Vector b)
|
||
|
|
asList f = V.fromList . f . V.toList
|
||
|
|
|
||
|
|
fromLists :: [[a]] -> Matrix a
|
||
|
|
fromLists = V.fromList . map V.fromList
|
||
|
|
|
||
|
|
|
||
|
|
gauss :: (Eq a, Fractional a) => Matrix a -> Vector a -> Vector a
|
||
|
|
gauss a b = x
|
||
|
|
where b' = fmap pure b
|
||
|
|
a' = V.zipWith (V.++) a b'
|
||
|
|
x = resubstitute (triangular a')
|
||
|
|
|
||
|
|
|
||
|
|
triangular :: (Eq a, Fractional a) => Matrix a -> Matrix a
|
||
|
|
triangular m
|
||
|
|
| null m = empty
|
||
|
|
| otherwise = cons r (triangular rs')
|
||
|
|
where (r, rs) = decons (rotatePivot m)
|
||
|
|
rs' = fmap f rs
|
||
|
|
f bs
|
||
|
|
| (V.head bs) == 0 = V.drop 1 bs
|
||
|
|
| otherwise = V.drop 1 $ V.zipWith (-) (fmap (*c) bs) r
|
||
|
|
where c = (V.head r)/(V.head bs)
|
||
|
|
|
||
|
|
|
||
|
|
rotatePivot :: (Eq a, Fractional a) => Matrix a -> Matrix a
|
||
|
|
rotatePivot m
|
||
|
|
| (V.head r) /= 0 = m
|
||
|
|
| otherwise = rotatePivot (snoc rs r)
|
||
|
|
where (r,rs) = decons m
|
||
|
|
|
||
|
|
|
||
|
|
resubstitute :: (Num a, Fractional a) => Matrix a -> Vector a
|
||
|
|
resubstitute = V.reverse . f . V.reverse . fmap V.reverse
|
||
|
|
where f m
|
||
|
|
| null m = empty
|
||
|
|
| otherwise = cons x (f rs')
|
||
|
|
where (r,rs) = decons m
|
||
|
|
x = (V.head r)/(V.last r)
|
||
|
|
rs' = fmap (asList subUnknown) rs
|
||
|
|
subUnknown (a1:(a2:as')) = ((a1-x*a2):as')
|