Go to Qiita Advent Calendar 2024 Top
LoginSignup
0
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@kei_tnk

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

Last updated at Posted at 2024-08-31

はじめに

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

行列の構造体を自作する

  • 当然、二次元配列そのままでも行列は表現できるけど、正方行列としての制約を組み込んだり、イニシャライザを楽に行えるように、構造体を新たに作ることとする。
  • scriptを利用して、i行、j列の値に「インスタンス名[i,j]」でアクセスできるようにしてみた。
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}
    }
}
  • 余り複雑なことしてないから、見れば分かって貰えるかな？
  • 実際に動かしてみると

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プロトコル」準拠としていたのに気付いただろうか？
  • このプロトコルは乗算が可能な型という制約を持つ。よって、この要素は乗算可能で有ることが担保される。
  • さて、乗算機能を追加するが、ついでに加減算機能も追加しておく。もしかしたら、将来役に立つかもしれないしね。
    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関数であることかな。
  • 一応､動作確認すると
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]

正方行列の累乗を高速に行う

  • これは、過去に投稿した繰り返し二乗法をつかうだけ。
  • だから、メソッドをpow()として、コードは以下の通り
    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)を適用できるようにしておこう
    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
    }
  • 一応、動作確認してみる
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構造体内で同じ演算子名「+=」であっても処理内容を変更する必要があるので、別投稿で動作確認を行った。
  • ついでに、要素がInt型の時のmod設定は、イニシャライズ時にのみ可能とした。また、mod計算は「+=」「-=」「*=」のときに、左辺のインスタンスのmodに従って行いこととし、普通の「+」「-」「*」では対応させていない。勿論、対応させられるけど、面倒だからサボりました。スミマセン。
  • でもって、できあがりは、以下の通り。
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」とか、使う演算子だけで良いかもね。まぁ、コードの長さは配点に関係ないから、何も考えず､全部コピーしても良いけどね。

おまけ

  • そういえば、できあがった正方行列の使い道って、ほぼ全て、列ベクトルに乗じることだよね。
  • なので、「*」の二項演算子に列ベクトルへの乗算を加えるね。
    //列ベクトルへの乗算
    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
    }
  • 計算例は以下のとおり。
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完できるように、覚えたアルゴリズムを自分のものにした方がレーティングは伸びるんだろうけど、新しいアルゴリズムを覚える方が楽しいんだよね。仕方ないね。
0
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?