swiftで正方行列計算【競技プログラミング】

# はじめに
- 競技プログラミングを解くテクニックの一つに、dp（動的計画法）があるけど、これを更に高速化するため、dpの内容を正方行列で表現し、遷移を正方行列の累乗で表現するという手法がある。
- だけど、swiftには行列計算の便利なライブラリはないっぽいので、自作することとする。
- ついでに、[累乗計算の高速化](https://qiita.com/kei_tnk/items/a622eec4e7a9d0c47aca)まで、実装してしまおうと思う。

# 行列の構造体を自作する
- 当然、二次元配列そのままでも行列は表現できるけど、正方行列としての制約を組み込んだり、イニシャライザを楽に行えるように、構造体を新たに作ることとする。
- [script](https://qiita.com/kei_tnk/items/693f777382b6b1e89eda)を利用して、i行、j列の値に「インスタンス名[i,j]」でアクセスできるようにしてみた。
```swift
struct Matrix<Element>:CustomStringConvertible where Element:Numeric{
    let size:Int // 正方行列のサイズ (size)×(size)の行列となる。
    var data:[[Element]] // 正方行列を、二次元配列として格納している。
    
    var description : String { // CustomStringConvertibleプロトコルとセットで、print内容を定義する
        var str = ""
        for raw in data {
            str += raw.description + "\n"
        }
        str.removeLast()
        return str
    }
    
    init(_ cnt:Int, _ v:Element = Element.zero){ // サイズと各要素の値v(初期値は0)を指定して正方行列を生成
        self.size = cnt
        self.data = [[Element]](repeating:[Element](repeating:v,count:cnt),count:cnt)
    }
    init(_ cnt:Int, _ vs:[Element]){ // サイズと各要素に格納する値の配列vsを指定して、正方行列を生成
        self.size = cnt
        self.data = [[Element]](repeating:[Element](repeating:vs[0],count:cnt),count:cnt)
        var num = 0
        OUT:for i in 0..<cnt {
            for j in 0..<cnt {
                data[i][j] = vs[num]
                num += 1
                if num > vs.count {break OUT}
            }
        } 
    }
       
    // i行目、j列目を表示（ただし、0起点）
    subscript(_ i:Int,_ j:Int) -> Element {
        get {data[i][j]}
        set {data[i][j] = newValue}
    }
}
```
- 余り複雑なことしてないから、見れば分かって貰えるかな？
- 実際に動かしてみると
```swift

var mat1 = Matrix(3,[1,2,3,4,5,6,7,8,9])
var mat2 = Matrix(4,[1,2,3,4,5])
print(mat1) // 出力①
print(mat2) // 出力②

mat2[3,1] = 100
print(mat2) // 出力③

///出力
①
[1, 2, 3]
[4, 5, 6]
[7, 8, 9]
②
[1, 2, 3, 4]
[5, 0, 0, 0]
[0, 0, 0, 0]
[0, 0, 0, 0]
③
[1, 2, 3, 4]
[5, 0, 0, 0]
[0, 0, 0, 0]
[0, 100, 0, 0]
```

# 正方行列に乗算の機能を追加する
- 実は、前記コードで、Matrix構造体の要素の型Elementを「Numericプロトコル」準拠としていたのに気付いただろうか？
- このプロトコルは乗算が可能な型という制約を持つ。よって、この要素は乗算可能で有ることが担保される。
- さて、乗算機能を追加するが、ついでに加減算機能も追加しておく。もしかしたら、将来役に立つかもしれないしね。
```swift
    static func +(left:Matrix,right:Matrix)->Matrix{
        let cnt = left.size
        var newMat = Matrix(cnt)
        for i in 0..<cnt {
            for j in 0..<cnt {
                newMat[i,j] = left[i,j] + right[i,j]
            }
        }
        return newMat
    }
    
    static func +=(left:inout Matrix,right:Matrix){
        let cnt = left.size
        for i in 0..<cnt {
            for j in 0..<cnt {
                left[i,j] = left[i,j] + right[i,j]
            }
        }
    }

    static func -(left:Matrix,right:Matrix)->Matrix{
        let cnt = left.size
        var newMat = Matrix(cnt)
        for i in 0..<cnt {
            for j in 0..<cnt {
                newMat[i,j] = left[i,j] - right[i,j]
            }
        }
        return newMat
    }
    
    static func -=(left:inout Matrix,right:Matrix){
        let cnt = left.size
        for i in 0..<cnt {
            for j in 0..<cnt {
                left[i,j] = left[i,j] - right[i,j]
            }
        }
    }
    
    static func *(left:Matrix,right:Matrix)->Matrix{
        let cnt = left.size
        var newMat = Matrix(cnt)
        for i in 0..<cnt {
            for j in 0..<cnt {
                for k in 0..<cnt {
                    newMat[i,j] += left[i,k] * right[k,j]                    
                }
            }
        }
        return newMat
    }
    
    static func *=(left:inout Matrix,right:Matrix){
        let cnt = left.size
        var newMat = Matrix(cnt)
        for i in 0..<cnt {
            for j in 0..<cnt {
                for k in 0..<cnt {
                    newMat[i,j] += left[i,k] * right[k,j]                    
                }
            }
        }
        left = newMat
    }
```
- まぁ、これも特に問題ないかな。ちょっとだけ特徴が有るとすれば、二項演算子を使うときは、引数名として、`left`と`right`が使われるということと、二項演算子の計算相手は外部のインスタンスだから、static関数であることかな。
- 一応､動作確認すると
```swift
var mat1 = Matrix(3,[1,2,3,4,5,6,7,8,9])
var mat2 = Matrix(3,[1,10,100,1,10,100,1,10,100])

print(mat1) // 出力①
print(mat2) // 出力②

mat1 *= mat2
print(mat1) // 出力③

///出力
①
[1, 2, 3]
[4, 5, 6]
[7, 8, 9]
②
[1, 10, 100]
[1, 10, 100]
[1, 10, 100]
③
[6, 60, 600]
[15, 150, 1500]
[24, 240, 2400]
```

# 正方行列の累乗を高速に行う
- これは、過去に投稿した[繰り返し二乗法](https://qiita.com/kei_tnk/items/a622eec4e7a9d0c47aca)をつかうだけ。
- だから、メソッドをpow()として、コードは以下の通り
```swift
    func pow(_ n:Int)->Matrix {
        var ans = self
        var x = self
        var i = n - 1
        while i > 0 {
            if i % 2 == 1 {ans *= x}
            x *= x
            i >>= 1
        }
        return ans
    }
```
- ついでに、要素にmod(_ m:Int)を適用できるようにしておこう
```swift
    func mod(_ m:Element)->Matrix where Element == Int {
        var ans = self
        for i in 0..<size {
            for j in 0..<size {
                ans[i,j] %= m
            }
        }
        return ans
    }
    
    func powm(_ n:Int,_ m:Int)->Matrix where Element == Int{
        var ans = self
        var x = self
        var i = n - 1
        while i > 0 {
            if i % 2 == 1 {ans *= x}
            x *= x
            i >>= 1
            
            ans = ans.mod(m)
            x = x.mod(m)
        }
        return ans
    }
```
- 一応、動作確認してみる
```swift
let mat_2 = Matrix(3,[2,0,0,0,2,0,0,0,2])
let mat_2_7 = mat_2.pow(7)
print(mat_2) // 出力①
print(mat_2_7) // 出力②

let mat_s = Matrix(3,[1,2,3,4,5,6,7,8,9])
let mat_s_2 = mat_s.pow(2)
print(mat_s) // 出力③
print(mat_s_2) // 出力④

print(mat_s_2.mod(10)) // 出力⑤
print(mat_s.powm(2,10)) // 出力⑥

let mat = Matrix(2,1.0)
print(mat) // 出力⑦
// print(mat.mod(10)) // エラー（Intしかmod出来ない）

///出力
①
[2, 0, 0]
[0, 2, 0]
[0, 0, 2]
②
[128, 0, 0]
[0, 128, 0]
[0, 0, 128]
③
[1, 2, 3]
[4, 5, 6]
[7, 8, 9]
④
[30, 36, 42]
[66, 81, 96]
[102, 126, 150]
⑤
[0, 6, 2]
[6, 1, 6]
[2, 6, 0]
⑥
[0, 6, 2]
[6, 1, 6]
[2, 6, 0]
⑦
[1.0, 1.0]
[1.0, 1.0]
```

# 反省
- 上記のmodの仕方では、毎回、「+=」計算後にmodを行っているが、「+=」計算内でオーバーフローしてしまうことがあるようで、累乗の回数nが大きいときに上手くいかないケースがあった。
- よって、「+=」の計算内でmod計算を行うこととした。
- なお、これは要素の型がInt型の時のみの対応であり、要素がDouble型の時はmodが使えないため、Matrix構造体内で同じ演算子名「+=」であっても処理内容を変更する必要があるので、[別投稿で動作確認を行った](https://qiita.com/kei_tnk/items/49dedb61f8f3d8c91628)。
- ついでに、要素がInt型の時のmod設定は、イニシャライズ時にのみ可能とした。また、mod計算は「+=」「-=」「\*=」のときに、左辺のインスタンスのmodに従って行いこととし、普通の「+」「-」「\*」では対応させていない。勿論、対応させられるけど、面倒だからサボりました。スミマセン。
- でもって、できあがりは、以下の通り。
```swift
struct Matrix<Element:Numeric>:CustomStringConvertible{
    let size:Int
    var data:[[Element]]
    let mod:Int?
    
    var description : String {
        var str = ""
        for raw in data {
            str += raw.description + "\n"
        }
        str.removeLast()
        return str
    }
    
    init(_ cnt:Int, _ v:Element = Element.zero){
        self.size = cnt
        self.data = [[Element]](repeating:[Element](repeating:v,count:cnt),count:cnt)
        self.mod = nil
    }
    init(_ cnt:Int, _ vs:[Element]){
        self.size = cnt
        self.data = [[Element]](repeating:[Element](repeating:Element.zero,count:cnt),count:cnt)
        self.mod = nil
        var num = 0
        OUT:for i in 0..<cnt {
            for j in 0..<cnt {
                data[i][j] = vs[num]
                num += 1
                if num == vs.count {break OUT}
            }
        } 
    }
    init(_ cnt:Int, _ v:Element = Element.zero, mod:Int? = nil) where Element == Int {
        self.size = cnt
        self.data = [[Element]](repeating:[Element](repeating:v,count:cnt),count:cnt)
        self.mod = mod
    }
    init(_ cnt:Int, _ vs:[Element],mod:Int? = nil) where Element == Int {
        self.size = cnt
        self.data = [[Element]](repeating:[Element](repeating:Element.zero,count:cnt),count:cnt)
        self.mod = mod
        var num = 0
        var vs_mod = vs
        if let m = mod {vs_mod = vs.map{$0 % m}}
        OUT:for i in 0..<cnt {
            for j in 0..<cnt {
                data[i][j] = vs_mod[num]
                num += 1
                if num == vs.count {break OUT}
            }
        } 
    }
    
    // i行目、j列目を表示（ただし、0起点）
    subscript(_ i:Int,_ j:Int) -> Element {
        get {data[i][j]}
        set {data[i][j] = newValue}
    }
    
    static func +(left:Self,right:Self)->Self{
        let cnt = left.size
        var newMat = Self(cnt)
        for i in 0..<cnt {
            for j in 0..<cnt {
                newMat[i,j] = left[i,j] + right[i,j]
            }
        }
        return newMat
    }

    static func +=(left:inout Self,right:Self){
        let cnt = left.size
        for i in 0..<cnt {
            for j in 0..<cnt {
                left[i,j] = left[i,j] + right[i,j]
            }
        }
    }    
    static func +=(left:inout Self,right:Self) where Element == Int {
        let cnt = left.size
        for i in 0..<cnt {
            for j in 0..<cnt {
                left[i,j] = left[i,j] + right[i,j]
                if let m = left.mod {left[i,j] %= m}
            }
        }
    }

    static func -(left:Self,right:Self)->Self{
        let cnt = left.size
        var newMat = Self(cnt)
        for i in 0..<cnt {
            for j in 0..<cnt {
                newMat[i,j] = left[i,j] - right[i,j]
            }
        }
        return newMat
    }
    
    static func -=(left:inout Self,right:Self){
        let cnt = left.size
        for i in 0..<cnt {
            for j in 0..<cnt {
                left[i,j] = left[i,j] - right[i,j]
            }
        }
    }
    static func -=(left:inout Self,right:Self) where Element == Int {
        let cnt = left.size
        for i in 0..<cnt {
            for j in 0..<cnt {
                left[i,j] = left[i,j] - right[i,j]
                if let m = left.mod {left[i,j] %= m}
            }
        }
    }
    
    static func *(left:Self,right:Self)->Self{
        let cnt = left.size
        var newMat = Self(cnt)
        for i in 0..<cnt {
            for j in 0..<cnt {
                for k in 0..<cnt {
                    newMat[i,j] += left[i,k] * right[k,j]
                }
            }
        }
        return newMat
    }
    
    static func *=(left:inout Self,right:Self){
        let cnt = left.size
        var newMat = Self(cnt)
        for i in 0..<cnt {
            for j in 0..<cnt {
                for k in 0..<cnt {
                    newMat[i,j] += left[i,k] * right[k,j]
                }
            }
        }
        left = newMat
    }
    static func *=(left:inout Self,right:Self) where Element == Int {
        let cnt = left.size
        var newMat = Self(cnt,mod:left.mod)
        for i in 0..<cnt {
            for j in 0..<cnt {
                for k in 0..<cnt {
                    newMat[i,j] += left[i,k] * right[k,j]
                    if let m = left.mod {newMat[i,j] %= m}
                }
            }
        }
        left = newMat
    }

    func pow(_ n:Int)->Self {
        var ans = self
        var x = self
        var i = n - 1
        while i > 0 {
            if i % 2 == 1 {ans *= x}
            x *= x
            i >>= 1
        }
        return ans
    }
    func pow(_ n:Int)->Self where Self == Matrix<Int> {
        var ans = self
        var x = self
        var i = n - 1
        while i > 0 {
            if i % 2 == 1 {ans *= x}
            x *= x
            i >>= 1
        }
        return ans
    }
    
    //列ベクトルへの乗算
    static func *(left:Self,right:[Element])->[Element]{
        let cnt = left.size
        var ans = [Element](repeating:Element.zero,count:cnt)
        for i in 0..<cnt {
                for j in 0..<cnt {
                    ans[i] += left[i,j] * right[j]                    
                }
        }
        return ans
    }

}

var mat = Matrix(3,[1,2,3,4,5,6,7,8,9],mod:5)
print(mat) // 出力①

mat += mat
print(mat) // 出力②

mat *= mat
print(mat) // 出力③

mat = mat.pow(1000_00000_00000_00000) // 10^18
print(mat) // 出力④

let vs = [10,20,30]
print(mat * vs) // 出力②
```
- 無駄に長いから、実際に競技プログラミングで使うときは、演算子は「*=」や「pow」とか、使う演算子だけで良いかもね。まぁ、コードの長さは配点に関係ないから、何も考えず､全部コピーしても良いけどね。


# おまけ
- そういえば、できあがった正方行列の使い道って、ほぼ全て、列ベクトルに乗じることだよね。
- なので、「*」の二項演算子に列ベクトルへの乗算を加えるね。
```swift
    //列ベクトルへの乗算
    static func *(left:Matrix,right:[Element])->[Element]{
        let cnt = left.size
        var ans = [Element](repeating:Element.zero,count:cnt)
        for i in 0..<cnt {
                for j in 0..<cnt {
                    ans[i] += left[i,j] * right[j]                    
                }
        }
        return ans
    }
```
- 計算例は以下のとおり。
```swift
let mat = Matrix(3,[1,2,3,4,5,6,7,8,9])
print(mat) // 出力①

let vs = [10,20,30]
print(mat * vs) // 出力②

///出力
①
[1, 2, 3]
[4, 5, 6]
[7, 8, 9]
②
[140, 320, 500]
```


# おわりに
- 今回の正方行列の自作で、競技プログラミングで戦う武器が増えるよ！やったね○○ちゃん！
- でも、高度なアルゴリズムを覚えても、使いこなせないと意味ないんだよね。
- 正直、新しいアルゴリズムを覚えるより、AtCoderのABCコンテストで4完できるように、覚えたアルゴリズムを自分のものにした方がレーティングは伸びるんだろうけど、新しいアルゴリズムを覚える方が楽しいんだよね。仕方ないね。