LoginSignup
40
34

More than 3 years have passed since last update.

Swift版 競プロ用チートシート(初心者向け)

Last updated at Posted at 2020-09-19

はじめに

こんにちは!屋代高校PC同好会結成会のTARDIGRADEです。
本記事では、Swiftを使って作成した競技プログラミング用チートシートをまとめています。あくまで「そのままコピペして使えるようなチートシートをまとめる」ということが目的ですので、アルゴリズムやデータ構造の仕組みを理解したいという方は、別の記事をご覧ください。一応リンクも貼ってあります。
便利そうなものが新しく見つかったら、随時更新していこうと思っているので、「こういうのが欲しいな~」という意見があればコメントください。また、「俺こんなの作ったよ!」というものがあれば、編集リクエストを送っていただければ積極的に採用したいと思います(一応内容の精査はするつもりです)。
Atcoderのコンテストでの提出数を見ても、Swiftで競技プログラミングをやってる方はまだまだ少ないですので、どんどん情報交換していきましょう!

・対象者

Swiftの基本文法(条件分岐、ループ文、関数の作成等)がわかっている方。
本記事では、アルゴリズムやデータ構造をメインに扱うつもりですので、簡単な標準ライブラリ(sort,abs,max等)や、計算量オーダーについての知識もあるほうが、わかりやすいかと思います。
また、タイトルに(初心者向け)とありますが、DFS以降の内容はAtCoder Beginner ContestのD問題以上で頻出するレベルのものであるということを先に断っておきます(ACLの登場で、今後C問題以下が難化する可能性は否めませんが…)。

・実行環境

本記事で作成したすべてのコードは、2020/10/27の時点で、AtCoderのジャッジシステム上で正確に動くことを確認しています。

・注意(必ず読んでください)

いまさら何言ってんだと思われるかもしれませんが、僕はこの記事を書くまで、Swiftにほとんど触れたことがありませんでした。じゃあなんで書こうと思ったんだって話なんですが、それはこの記事の主旨から外れるので省略させてください。一言でいえば友達のためです。
完全独学で経験値もほぼなし、細かい文法なんかは全然抑えられてないです。そのため、普段からSwiftを使って作業されている方から見ると、無駄に長かったり、読みにくかったり、違和感のある書き方をしてしまっているかもしれません。目に余るようでしたら修正しますので、遠慮なく指摘してください。
また、僕は緑コーダーです。難しいアルゴリズムやデータ構造は理解できていない部分もあります。そういう意味で、タイトルにもある通りこの記事は初心者向けということにさせていただいています。
以上の点を理解したうえで、本記事を参考にしたり、あるいは引用していただけるとありがたいです。

標準入力の受け取り方

クリックで展開

標準入力は、1行ずつ受け取るのがわかりやすくていいと思います。

// 入力値が一つで、String型として扱いたいとき
// "abc" -> "abc"
let input = readLine()!

// 入力値が一つで、Int型として扱いたいとき
// "6" -> 6
let input = Int(readLine()!)!

// 入力値が複数あり(スペースで区切られている)、String型として扱いたいとき
// "abc" "def" -> ["abc","def"]
let input = readLine()!.split(separator: " ")

// 入力値が複数あり(スペースで区切られている)、Int型として扱いたいとき
// "123" "456" -> [123,456]
let input = readLine()!.split(separator:" ").map{Int($0)!}

// 入力値が一つでString型であるものを、配列として扱いたいとき
// "abcdef" -> ["a","b","c","d","e","f"]
let input = Array(readLine()!).map{String($0)}

mapはとても便利な関数ですので、知らなかった方はぜひ調べてみてください。こちらの記事でも取り上げられています。

また、文字列の受け取りに関してですが、Swiftでとある文字列に対して「指定した位置の文字を取得する」という操作をしようとすると、少々面倒な書き方をしなければなりません(もしかしたらすっきり書く方法もあるかも…)。ですので、入力された文字列に対して何か操作を行う場合は、一番下にあるように文字列を分割して配列に保存するとよいと思います。

素数列挙・素数判定

2020/10/12 追記
素数列挙・判定のプログラムを、より高速に動くものに更新しました。
プログラムを提供してくださった@En3_HCl様には心より感謝申し上げます。

2020/10/16 追記
素数列挙のプログラムのバクを修正&高速化して更新しました。
プログラムを提供してくださった@koher様には心より感謝申し上げます。

クリックで展開

「エラトステネスのふるい」を利用することで、$n$以下の素数の列挙を高速に行うことができます。エラトステネスのふるいについてはこちら

作成したのは、$n$以下の素数の列挙する関数です。例えば、$10$を渡すと$[2,3,5,7]$が返ってきます。

func primes(upTo number: Int) -> [Int] {
    precondition(number >= 0)
    if number < 2 { return [] }
    var sieve: [Bool] = .init(repeating: false, count: number + 1)
    for m in stride(from: 3, through: Int(Double(number).squareRoot() + 1.5), by: 2) {
        if sieve[m] { continue }
        let maxK = number / m
        if maxK < 2 { continue }
        for k in 2 ... maxK {
            sieve[k * m] = true
        }
    }
    var result: [Int] = [2]
    for m in stride(from: 3, through: number, by: 2) {
        if sieve[m] { continue }
        result.append(m)
    }
    return result
}

ざっくり時間を測ってみた感じ、$n<=10^7$までならかなり余裕をもって、$n=10^8$でもなんとか2秒以内に終わります。

ついでに、自然数$n$が素数かどうかを判定する関数も作成しました。こっちはただ割っていくだけなのでとてもシンプルです。計算量も$O(√n)$とかなり高速です。

func check_primes(n: Int) -> Bool {
    if n == 2{
        return true
    }
    if n.isMultiple(of: 2){
        return false
    }
    // 計算するのは√nまでで十分
    // 2の倍数をスキップして計算する
    for i in stride(from: 3, to: min(Int(sqrt(Double(n)))+1, n), by: 2){
        if n.isMultiple(of: i){
            return false
        }
    }
    return true
}

実際に問題を解く中で上記のアルゴリズムを使ってみたい方は、下の問題を解いてみてください。基本問題です。
AtCoder Beginner Contest 149 C - Next Prime
Swiftでの解答例はこちら(詳しい解法は公式の解説をご覧ください)。

素因数分解

クリックで展開

素因数分解を行うアルゴリズムはいくつか知られています。詳しくはこちらの記事を読んでください(正直僕は半分くらいしか理解できてないです…)。

本記事には、試し割り法と呼ばれるアルゴリズムのチートシートを載せようと思います。一つの自然数に対して素因数分解を行うのであれば、この方法で十分高速に計算ができるからです。

import Foundation
// 第一引数に、素因数分解したい自然数を渡す
func prime_factorization(n:Int)->[[Int]]{
    // n = 1 なら素因数分解できないので空の二次元配列を返す
    if n == 1{
        return []
    }
    // 戻り値となる二次元配列を初期化
    var arr:[[Int]] = []
    // nの値を変数に格納
    var temp = n
    // 計算するのは√nまでで十分
    for i in 2...Int(sqrt(Double(n)))+1{
        // 割り切れる -> その数を素因数に持つ
        if temp%i==0{
            var cnt=0
            // 素因数の字数を計算
            while temp%i==0{
                // 割り切れるなら次数を一つ増やす
                cnt+=1
                // tempをiで割る
                temp /= i
            }
            // 素因数とその次数をarrに追加
            arr.append([i,cnt])
        }
    }
    // temp ≠ 1ならtempは素因数なのでarrに追加
    if temp != 1{
        arr.append([temp, 1])
    }
    // arrが空なら n = (素数) なのでarrに追加
    if arr==[] {
        arr.append([n, 1])
    }
    // arrを返す
    return arr
}

複数の自然数に対して素因数分解を行う場合は、エラトステネスのふるいの考え方を利用したアルゴリズムを用いることで、より高速に計算できるようになります。興味のある方はこちらの記事を読んでください。

要望があれば(あるいはもう少し時間に余裕ができたら)、こちらのアルゴリズムのチートシートも作って載せるつもりではいますが、素数列挙の時に使ったエラトステネスのふるいのアルゴリズムと大差ない物なので、とりあえず今は省略させていただきます(「私作れました!」という方がいましたら、編集リクエストを送っていただければ積極的に採用させていただきます)。

実際に問題を解く中で素因数分解を使ってみたい方は、下の問題を解いてみてください。「The素因数分解!」という感じの問題です。
AtCoder Beginner Contest 169 D - Div Game
Swiftでの解答例はこちら(詳しい解法は公式の解説をご覧ください)。

約数列挙

クリックで展開

約数列挙は素数列挙などに比べてだいぶシンプルです。なんせ割っていくだけなので…。アルゴリズムの考え方はこちらの記事を読んでください。
import Foundation
// 第一引数には約数を列挙したい自然数を渡す
func print_divisors(n:Int)->[Int]{
    // 戻り値となる配列の初期化
    var divisors:[Int] = []
    // ループを回すのは√nまでで十分
    for i in 1...Int(sqrt(Double(n))){
        // 割り切れるならiはnの約数
        if n % i == 0{
            divisors.append(i)
            // i != n / i なら i^2 != n
            if i != n / i{
                divisors.append(n/i)
            }
        }
    }
    // 配列をソートする
    divisors.sort{ $0 < $1 }
    // divisorsを返す
    return divisors
}

最後にsortしているので、計算量は$O(nlogn)$でしょうか(あんまり自信ないです、すみません…)。sortの部分をコメントアウトすれば計算量は$O(√n)$になります。

実際に問題を解く中で約数列挙を使ってみたい方は、下の問題を解いてみてください。一瞬「えっ?」となるかもしれませんが、ちょっとした考察を入れるとただの約数列挙の問題に帰着します。
AtCoder Beginner Contest 144 C - Walk on Multiplication Table
Swiftでの解答例はこちら(詳しい解法は公式の解説をご覧ください)。

bit全探索

クリックで展開

全探索についてはチートシートを書くかどうか迷いましたが、本記事では多重ループでは上手くいかないものについてのみ扱おうかと思います(要望があればノーマルな全探索についてもまとめます)。

bit全探索は$n$個の要素からなる集合のすべての部分集合を調べ上げる手法のことです。詳しくはこちらの記事を読んでください。また、bit全探索にはbit演算子等の知識も必要になります。こちらにわかりやすくまとめてあったので、不安な方は一度ご確認ください。

作成したチートシートについてですが、関数化はしていません。関数の形にすると逆にわかりにくくなる気がしたので…。

import Foundation
// 要素数nの集合についてbit全探索を行う
// 0~n^2のループ
for bit in 0..<(1<<n){
    // 順に左にシフトさせ最下位bitのチェックを行う
    for i in 0..<n{
        // フラグが立っているかどうかの判定
        if bit & 1<<i > 0 {
            // フラグが立っていた時の処理
        }
        else{
            // フラグが立っていなかった時の処理
        }
    }
}

これを使用する際には、$n$に要素数を入れることと、問題に合わせて処理の中身をきちんと記述することを忘れないようにしてください。

実際に問題を解く中でbit全探索を使ってみたい方は、下の問題を解いてみてください。少し問題が複雑ですが、いろいろと考察して解いてみることで、bit全探索の理解が深まると思います。
AtCoder Beginner Contest 128 C - Switches
Swiftでの解答例はこちら(詳しい解法は公式の解説をご覧ください)。

順列全探索

2020/10/11 追記
プログラムを、より高速に動くものに更新しました。
プログラムを提供してくださった@koher様には心より感謝申し上げます。

クリックで展開

順列全探索に関しては、「これ!」というような記事が見つからなかったので、軽くですが僕から説明します。わかりにくければグーグル先生を頼るか、コメントで質問してください。

順列全探索とは、一口で言えば$n!$通りの全探索のことです。例えば、$(1,2,3)$という配列がある場合、この配列について順列全探索を行うとは、$(1,2,3),(1,3,2),(2,1,3),(2,3,1),(3,1,2),(3,2,1)$の6種類の並び方について探索することを指します。

PythonやC++では順列を生成してくれる標準ライブラリがあったりするんですが、Swiftにはなさそうなので、自作するしかありません。もしそのような標準ライブラリがあれば、本記事と比べてかなりシンプルに順列全探索を書けるようになるはずですので、ご存じの方がいれば教えてください。

本記事で作成したのは、$ N$個の要素を持つ配列$a$が与えられているとき、$a[0]~ a[N-1]$の要素を並べ替えて得られる$ N!$個の順列を、二次元配列の形でまとめて返す関数です。つまり、$(1,2,3)$を渡すと、$[(1,2,3),(1,3,2),(2,1,3),(2,3,1),(3,1,2),(3,2,1)]$が返ってきます。ただし、sortがされているとは限らないのでその点は注意してください。

extension Sequence {
    func permutations() -> [[Element]] {
        func _permutations<T>(of values: [T], indices: Range<Int>, result: inout [[T]]) {
            if indices.isEmpty {
                result.append(values)
                return
            }
            var values = values
            for i in indices {
                values.swapAt(indices.lowerBound, i)
                _permutations(of: values, indices: (indices.lowerBound + 1) ..< indices.upperBound, result: &result)
            }
        }

        var result: [[Element]] = []
        let values = Array(self)
        _permutations(of: values, indices: values.indices, result: &result)
        return result
    }
}

実際に問題を解く中で順列全探索を使ってみたい方は、下の問題を解いてみてください。問題に合わせて若干関数の形を変えるか、あるいは使い方を工夫する必要がありますが、順列全探索の基本問題です。
AtCoder Beginner Contest 145 C - Average Length
Swiftでの解答例はこちら(詳しい解法は公式の解説をご覧ください)。

二分探索

2020/10/11 追記
プログラムを、より一般的に動くものに更新しました。具体的な変更点は以下の二点です。
・Int 以外の要素を持つ Array に対応
・Array 以外のコレクションに対応
プログラムを提供してくださった@koher様には心より感謝申し上げます。

クリックで展開

二分探索って何?という方はこちらの記事を読んでください。

本記事で作成したものは、正確に言えば「ソート済みの配列aのイテレータのうち、検索したい値以上となる最小のイテレータを返す関数」です。

func values<C>(greaterThanOrEqualTo value: C.Element, in sortedCollection: C) -> C.SubSequence
        where C: RandomAccessCollection, C.Element: Comparable, C.Index == Int {
    var low = sortedCollection.startIndex - 1
    var high = sortedCollection.endIndex
    while high - low > 1 {
        let mid = low + (high - low) / 2
        if sortedCollection[mid] >= value {
            high = mid
        } else {
            low = mid
        }
    }
    return sortedCollection[high...]
}

使い方は↓

values(greaterThanOrEqualTo: 3, in: [2, 2, 3, 3, 3, 4, 4, 4, 4]) // [3, 3, 3, 4, 4, 4, 4]
values(greaterThanOrEqualTo: "e", in: ["a", "b", "c", "d", "e", "f", "g"]) // ["e", "f", "g"]
values(greaterThanOrEqualTo: 7, in: [2, 3, 5, 7, 11]).startIndex // 3
values(greaterThanOrEqualTo: 7, in: [2, 3, 5, 7, 11]).indices // 3..<5

実際に問題を解く中で二分探索を使ってみたい方は、下の問題を解いてみてください。少し考察が必要で、難易度的にもちょうどいいかと思います。
AtCoder Beginner Contest 077 C - Snuke Festival
Swiftでの解答例はこちら(詳しい解法は公式の解説をご覧ください)。

スタック・DFS(深さ優先探索)

クリックで展開

DFSとは木やグラフを探索するためのアルゴリズムの一つです。
この章を読んでいる方で、「DFSが何なのか全くわからない」という方は少ないとは思いますが、一応いろいろとリンクを貼っておきます。

グラフと木
深さ優先探索(Depth First Search)の基本
スタックとキューを極める! 〜 考え方と使い所を特集 〜

本記事ではスタックver.のDFSのチートシートを載せます。再帰関数ver.はうまくチートシートの形に纏める自信がないのでここには載せませんが、Swiftの場合スタックよりも再帰のほうが断然シンプルに書けるので、興味がある方はぜひ調べてみてください。

では早速DFSのチートシートを載せます………と行きたいところだったんですが、Swiftにはスタックがないそうなので、自分で実装するか別のデータ構造で代用しなければなりません。いろいろと探してみたところ、既に作成している方がいましたので、そちらをほぼそのままコピペさせていただきます。オリジナルはこちらから。
ちなみに、明示的にスタックというものを宣言したい場合を除いて、これを使う必要はありません。実際中身はただのArray型の下位互換ですので、Array型で代用できます。Swiftにおけるスタック、キューについて教えてくださった@koher様には心より感謝申し上げます。

struct Stack<T> {
  var array = [T]()
  var isEmpty: Bool {
    return array.isEmpty
  }
  var count: Int {
    return array.count
  }
  mutating func push(_ element: T) {
    array.append(element)
  }
  mutating func pop() -> T? {
    return array.popLast()
  }
  var top: T? {
    return array.last
  }
  func makeIterator() -> AnyIterator<T> {
    var curr = self
    return AnyIterator {
      return curr.pop()
    }
  }
}
// スタックの生成
var a = Stack<Any>()
// スタックが空かどうかの判定
a.isEmpty
// スタックのサイズを返す
a.count
// オブジェクトをスタックに追加する
a.push(1)
// 一番上のオブジェクトを削除し、それを返す
a.pop()
// 一番上のオブジェクトを返す
a.top
// イテレータの作成
a.makeIterator()

Swift初心者の僕にとっては理解が難しい部分も多かったのですが、動くのでいったん良しとします。何か不備に気付いた方がいらっしゃいましたら、教えていただけるとありがたいです。

このスタックを使ってDFSの大枠を作ると、こんな感じになります。

import Foundation
struct Stack<T> {
  // 中身省略
}
// 第一引数には隣接リストで表したグラフを渡す
// 第二引数には探索をスタートする頂点を渡す
func dfs(graph:inout [[Int]], start:Int)->[Bool]{
    // スタック
    var s = Stack<Any>()
    // スタート地点の情報をプッシュ
    s.push([[0],graph[start]]) // (<原点から頂点iまでの距離,原点に隣接している頂点>)
    // 訪問済みの頂点を記録する
    var visited:[Bool] = []
    for _ in 0..<graph.count{
        visited.append(false)
    }
    visited[start] = true
    // スタックが空になるまでループを回す
    while !s.isEmpty{
        // ダウンキャストする
        let label = s.pop() as! [[Int]]
        // 今いる頂点とつながっている頂点についてループを回す
        for i in label[1]{
            // 探索済みであればスタックには追加しない
            if visited[i]{
                continue
            }
            // 未探索であればスタックに追加する
            else{
                visited[i] = true
                s.push([label[0][0]+1,graph[i]])
            }
        }
    }
    return visited
}

重ねて言いますが、これはざっくりとした大枠を書いただけですので、実際に使う際には問題に合わせてプログラムを書き加えたり書き換えたりしてください。

実際に問題を解く中でDFSを使ってみたい方は、下の問題を解いてみてください。緑diffですが、考察ができればこっちのもんです。ちなみにスタックver.でやろうとすると、慣れていないと実装が大変ですので、できる方は再帰関数を使って書くことをお勧めします。
AtCoder Beginner Contest 138 D - Ki
Swiftでの解答例はこちら(詳しい解法は公式の解説をご覧ください)。 → スタックver. / 再帰ver.

いつも解答例には、どのような考え方でプログラムを書いているのかが、ある程度わかるようにコメントをつけておりますが、今回はちょっとコメントが雑です。わかりやすく書けなかったのは僕の力不足です、すみません。
そもそも、記事に張り付けている問題や解答例にどれほど需要があるのか定かではないのですが、もし、「もう少しわかりやすく説明してほしい」という方がいれば、コメントあるいはTwitterのほうで質問をいただければ返したいと思います。

キュー・BFS(幅優先探索)

クリックで展開

BFSとは木やグラフを探索するためのアルゴリズムの一つです。
この章を読んでいる方で、「BFSが何なのか全くわからない」という方は少ないとは思いますが、一応いろいろとリンクを貼っておきます。

グラフと木
幅優先探索(Breadth First Search)の基本
スタックとキューを極める! 〜 考え方と使い所を特集 〜

BFSを実装する場合はキューというデータ構造が必要になります。ですが、例のごとくSwiftにキューはないので、自分で実装するか別のデータ構造で代用することになります。いろいろ調べたところ、リングバッファと呼ばれるデータ構造が使えそうだったので、これをベースにキューを実装してみたいと思います。
ちなみに、ArraySliceをキューの代わりに使うこともできます。メモリの消費がリングバッファよりも多いですが、それが競技プログラミングで問題になることはほぼないと思われます。
Swiftにおけるスタック、キューについて教えてくださった@koher様には心より感謝申し上げます。

struct Queue<T> {
  var array = [T]()
  var MAX:Int
  init(MAX: Int) {
    self.MAX = MAX
    array = Array(repeating: 0 as! T, count: MAX)
  }
  var head = 0
  var tail = 0
  var isEmpty: Bool {
    return head == tail
  }
  var isFull: Bool{
    return head == (tail + 1) % MAX
  }
  var count: Int {
    if head <= tail{
        return tail-head
    }
    else{
        return MAX-head+tail
    }
  }
  var bottom: T {
    return array[head]
  }
  mutating func enqueue(_ element: T) -> Bool {
    if head == (tail + 1) % MAX{
      return false
    }
    array[tail] = element
    if tail + 1 == MAX{
        tail = 0
    }
    else{
        tail += 1
    }
    return true
  }
  mutating func dequeue() -> T? {
    if head == tail{
        return nil
    }
    let x = array[head]
    if head + 1 == MAX{
        head = 0
    }
    else{
        head += 1
    }
    return x
  }
  mutating func print() -> ArraySlice<T>{
      return array[head..<tail]
  }
}
// MAXを指定してキューを生成
var a = Queue<Any>(MAX:100000)
// キューが空かどうかの判定
a.isEmpty
// キューが満杯かどうかの判定
a.isFull
// キューのサイズを返す
a.count
// 一番下のオブジェクトを返す
a.bottom
// オブジェクトをキューに追加する
a.enqueue(1)
// 一番下のオブジェクトを削除し、それを返す
a.dequeue()
// キューの中身を返す
a.print()

僕自身、リングバッファというものをこの機会に初めて学んだので、どこかに間違いがあるかもしれません。一応テストはしていますが、なにか気づいた方がいたら教えていただけるとありがたいです。

このキューを使ってBFSの大枠を作ると、こんな感じになります。

import Foundation
struct Queue<T> {
  // 中身省略
}
// 第一引数には隣接リストで表したグラフを渡す
// 第二引数には探索をスタートする頂点を渡す
// 戻り値はスタート地点から各頂点までの距離(到達不可能な場合は-1)
func bfs(g:inout [[Int]], start:Int)->[Int]{
    // 訪問済みの頂点を記録する
    var seen = Array(repeating:-1,count:n)
    seen[start] = 0
    // キュー
    var q = Queue<Any>(MAX: 100000)
    // スタート地点の情報をキューに追加
    q.enqueue(start)
    // キューが空になるまでループを回す
    while !q.isEmpty{
        let label = q.dequeue() as! Int
        // 今いる頂点とつながっている頂点についてループを回す
        for i in g[label]{
            // 未探索であればキューに追加する
            if seen[i] == -1{
                seen[i] = seen[label] + 1
                q.enqueue(i)
            }
        }
    }
    return seen
}

重ねて言いますが、これはざっくりとした大枠を書いただけですので、実際に使う際には問題に合わせてプログラムを書き加えたり書き換えたりしてください。

実際に問題を解く中でBFSを使ってみたい方は、下の問題を解いてみてください。今回はちょっと難しめかもしれません。脳筋でゴリゴリ計算していくと、TLEする可能性が高いので気を付けてください。
AtCoder Beginner Contest 176 D - Wizard in Maze
Swiftでの解答例はこちら(詳しい解法は公式の解説をご覧ください)。

プライオリティーキュー・ダイクストラ法

2020/10/27 追記
プライオリティーキューとダイクストラ法の実装を改良して更新しました。
プログラムを提供してくださった@koher様には心より感謝申し上げます。

また、更新前のプライオリティーキューは、別記事で@conf8o様が書かれていたものを引用させていただいておりました。改めて感謝申し上げます。
引用させていただいた記事:競プロ用優先度付きキュー(Priority Queue) in Swift

クリックで展開

ダイクストラ法とは、最短経路問題を解くためのアルゴリズムの一つです。負の重みをもつ辺が無い場合に使うことができます。詳しい原理についてはこちらのサイトを読んでください。

ダイクストラ法を知っている方、あるいは上で紹介した記事を読まれた方はわかるかと思いますが、ダイクストラ法を効率よく動かすためにはプライオリティーキュー(優先度付きキュー)と呼ばれるデータ構造が必要になります。
プライオリティーキューとは、「値の追加」と「最小の値を取り出し削除する」という操作を$O(logn)$で行うことができるデータ構造です。値が固定の場合は配列をソートするだけで同じようなことができますが、値が出たり入ったりする場合にはプライオリティーキューを使ったほうが高速です。

ご多分に漏れず、Swiftにはプライオリティーキューが無いので自力で実装する必要があります。自力での実装が厳しそうだったのでいろいろ検索したところ、Swiftでのプライオリティーキューの実装方法をまとめてくださっているQiitaの記事が見つかりました。今回はそちらのコードを引用させていただきます。

struct PriorityQueue<Element> {
    private var elements: [Element] = []
    private let areInIncreasingOrder: (Element, Element) -> Bool

    init<S>(_ elements: S, by areInIncreasingOrder: @escaping (Element, Element) -> Bool) where S: Sequence, S.Element == Element {
        self.areInIncreasingOrder = areInIncreasingOrder
        for element in elements {
            append(element)
        }
    }

    init(by areInIncreasingOrder: @escaping (Element, Element) -> Bool) {
        self.init(EmptyCollection(), by: areInIncreasingOrder)
    }

    var isEmpty: Bool { elements.isEmpty }
    var count: Int { elements.count }
    var first: Element? { elements.first }

    mutating func append(_ element: Element) {
        var i = elements.count
        elements.append(element)
        elements.withUnsafeMutableBufferPointer { elements in
            while i > 0 {
                let parentIndex = (i - 1) >> 1
                let parent = elements[parentIndex]
                guard areInIncreasingOrder(element, parent) else { break }
                elements[parentIndex] = element
                elements[i] = parent
                i = parentIndex
            }
        }
    }

    mutating func popFirst() -> Element? {
        guard let element = elements.popLast() else { return nil }
        guard let first = elements.first else { return element }

        elements.withUnsafeMutableBufferPointer { elements in
            elements[0] = element

            var  i = 0
            while true {
                var childIndex: Int = (i << 1) + 1
                guard childIndex < elements.count else { break }
                var child: Element = elements[childIndex]
                let rightIndex = childIndex + 1
                if rightIndex < elements.count {
                    let right = elements[rightIndex]
                    if areInIncreasingOrder(right, child) {
                        childIndex = rightIndex
                        child = right
                    }
                }
                if areInIncreasingOrder(element, child) { break }
                elements[childIndex] = element
                elements[i] = child
                i = childIndex
            }
        }

        return first
    }
}
extension PriorityQueue where Element: Comparable {
    init<S>(_ elements: S) where S: Sequence, S.Element == Element {
        self.init(elements, by: <)
    }

    init() {
        self.init(by: <)
    }

}

こちらを使って、以下のようにダイクストラ法を実装しました。

func dijkstra<Distance>(graph: [[(index: Int, distance: Distance)]], startedAt start: Int) -> [Distance?] where Distance: Comparable, Distance: AdditiveArithmetic {
    var result: [Distance?] = .init(repeating: nil, count: graph.count)
    result[start] = .zero
    var queue = _PriorityQueue<(Distance, Int)>(by: { $0.0 < $1.0 })
    queue.append((.zero, start))
    while let (totalDistanceToI, i) = queue.popFirst() {
        guard let minTotalDistanceToI = result[i] else { preconditionFailure("Never reaches here.") }
        if minTotalDistanceToI < totalDistanceToI { continue }
        assert(totalDistanceToI == minTotalDistanceToI)
        for (j, distance) in graph[i] {
            if let totalDistanceToJ = result[j] {
                if totalDistanceToI + distance < totalDistanceToJ {
                    let newTotalDistanceToJ = totalDistanceToI + distance
                    result[j] = newTotalDistanceToJ
                    queue.append((newTotalDistanceToJ, j))
                }
            } else {
                let newTotalDistanceToJ = totalDistanceToI + distance
                result[j] = newTotalDistanceToJ
                queue.append((newTotalDistanceToJ, j))
            }
        }
    }
    return result
}

実際に問題を解く中でダイクストラ法を使ってみたい方は、下の問題を解いてみてください。
かなり難しいです。茶色の方にとっては厳しいかもしれません。
AtCoder Beginner Contest 035 D - トレジャーハント
Swiftでの解答例はこちら(詳しい解法は公式の解説をご覧ください)。

UnionFind

クリックで展開

UnionFindは競プロで頻出するデータ構造の一つです。詳しい仕組みや実装の中身についてはこちらのサイトをご覧ください。

実装は、これまで書いてきたリングバッファやプライオリティーキューに比べると軽めです。

struct UnionFind {
    var par:[Int]
    var rank:[Int]
    var n:Int
    //最初は全てが根であるとして初期化
    init(n:Int){
        self.n = n
        self.par = Array(repeating:0,count:n)
        self.rank = Array(repeating:0,count:n)
        for i in 0..<n{
            par[i] = i
        }
    }
    // データxが属する木の根を再帰で得る:root(x) = {xの木の根}
    mutating func root(x:Int) -> Int {
        if par[x] == x{
            return x
        }
        else{
            return root(x:par[x])
        }
    }
    // xとyの木を併合
    mutating func unite (x:Int,y:Int){
        let x = root(x:x)
        let y = root(x:y)
        if x == y{
            return
        }
        if rank[x] < rank[y]{
            par[x] = y
        } else {
            par[y] = x
            if rank[x] == rank[y]{
                rank[x] += 1
            }
        }
    }
    // 2つのデータx, yが属する木が同じならtrueを返す
    mutating func same(x:Int,y:Int)->Bool{
        return root(x:x) == root(x:y)
    }
}

実際に問題を解く中でUnionFindを使ってみたい方は、下の問題を解いてみてください。UnionFindを使う典型的な問題で、当時でこそ緑上位diffくらいありますが、ACLがでた今だと恐らく灰Diffです。
AtCoder Regular Contest 032 B - 道路工事
Swiftでの解答例はこちら(詳しい解法は公式の解説をご覧ください)。

組み合わせ(nCk)

クリックで展開

競プロの問題を解いていると、「$nCk$ $mod$ $p$」を高速に計算したい場面にしばしば遭遇します。計算の方法はいろいろありますが、計算速度の観点から、こちらの記事で紹介されている方法を使うことが多いです。
紹介されているプログラムを、Swiftで書き直すとこんな感じです。
let MAX = 500000
let MOD = 1000000007
var fac = Array(repeating:0,count:MAX)
var finv = Array(repeating:0,count:MAX)
var inv = Array(repeating:0,count:MAX)
// テーブルを作る前処理
func COMinit(){
    fac[0] = 1
    fac[1] = 1
    finv[0] = 1
    finv[1] = 1
    inv[1] = 1
    for i in 2..<MAX{
        fac[i] = fac[i-1] * i % MOD
        inv[i] = MOD - inv[MOD%i] * (MOD / i) % MOD
        finv[i] = finv[i - 1] * inv[i] % MOD
    }
}

// 二項係数計算
func COM(n:Int,k:Int)->Int{
    if n < k {
        return 0
    }
    if n < 0 || k < 0{
        return 0
    }
    return fac[n] * (finv[k] * finv[n - k] & MOD) % MOD
}

// 前計算
COMinit()
// nCk
COM(n:n,k:k)

計算量は、前計算$O(MAX)$、クエリ処理$O(1)$です。
また、使用条件は、「$MOD$は素数」「$1<k<n<MAX$」です。

実際に問題を解く中でこのチートシートを使ってみたい方は、下の問題を解いてみてください。若干の考察が必要ですが、実装はチートシートを使えばかなり簡単にできます。
AtCoder Beginner Contest 145 D - Knight
Swiftでの解答例はこちら(詳しい解法は公式の解説をご覧ください)。

追加予定のチートシート

特になし。
セグ木やBIT、最大流などのより難しいデータ構造・アルゴリズムは、別で経験者向けのチートシートの記事を用意し、そちらで書くつもりです。
追加してほしいチートシートがございましたら、コメントもしくはTwitterのほうで受け付けております。

40
34
19

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
40
34