Zenn
🦜

AtCoderをSQLで解いて提出(ABC 170 E - Smart Infants)

2021/08/25に公開
Python
SQL
AtCoder
mutliset
abc170
tech

はじめに

AtCoderのサポート言語にSQLはありませんが、
Python標準搭載のSQLiteを用いるとSQLで解くことができます。
PythonユーザーにとってはAtCoderで稀に要求されるMultiSetというデータ構造の代替になりえます。

お題

ABC 170 E - Smart Infants

概要

次のような設定の問題です

  • 複数の幼稚園児と複数の幼稚園があります。
  • 幼稚園児はどこか1つの幼稚園に所属します。
  • 幼稚園はAtCoderのレートを持っています。(ぶっとんでる🤣)
  • 1人の幼稚園児の転園を行うクエリがQ回きます
  • クエリがくるごとに、転園後の「"各幼稚園の所属園児の最大レート"の最小値」を出力する問題です

通常解法

公式解説されているので詳細は割愛します。
ポイントは、次になります。

  • MultiSetというデータ構造が要求される
    • 最大・最小の取得、データの追加・削除がO(log N)で可能
  • MultiSetはPythonにはない
    • MultiSetの独自実装はすごく大変
    • 今回の問題については、ならし計算量とheapqを使うと回避することもできます。DNA1980さんの解説

SQLで解いてみる

SQLite3を使います
SQLite3のドキュメントはこちらで、
提出コードはこちらです

ポイント0: ローカル関数で高速化

def main():
  ...
  
main()
  • Pythonの定番テクニックです
    • グローバル変数へのアクセスが遅いので、全部をローカルに書くことで高速化します
  • PyPyだと大きな差はないことや、逆に遅くなるケースもありますが、一応入れます

ポイント1: インメモリDB

con = sqlite3.connect(':memory:', isolation_level=None)

ファイルに書き込むと遅いのでインメモリDBモードでSQLiteを起動します。

PythonのSQLite3は自動でコミットしたりする機能がありますが、isolation_level=Noneで無効化しておきます

ポイント2: 不要機能のOFF

PRAGMA trusted_schema = OFF;
PRAGMA journal_mode = OFF;
PRAGMA synchronous = OFF;
PRAGMA temp_store = memory;
PRAGMA secure_delete = OFF;

不要機能をOFFにして高速化します。
詳細はSQLiteのドキュメントにあります

ポイント3: テーブルとインデックス作成

CREATE TABLE infants (kindergarten INT, rate INT);
CREATE INDEX infant_index ON infants(kindergarten, rate);
CREATE TABLE kindergartens (rate INT);
CREATE INDEX rates ON kindergartens(rate);
CREATE TABLE ans(min_rate INT);
  • infants(園児)テーブルは各園児の所属幼稚園(kindergarten)とレートのテーブルです
  • kindergartens(幼稚園)テーブルは所属園児の最大レートのテーブルです
  • ansテーブルはクエリごとにその時点でのkindergartensの最小値を格納するテーブルです
    • クエリごとには答えを出力せずにansテーブルに答えを貯めておいて、最後にansテーブルの中身をまとめて出力して問題に解答します
    • Python <-> SQLiteのやりとりを減らして高速化できます

後に、SELECT MAX(rate) FROM infants WHERE kindergarten = ..というクエリを使うのですが、これを高速にできるようなインデックスを作ります
それが CREATE INDEX infant_index ON infants(kindergarten, rate);の部分です。

  • インデックスのMAXは2分探索で高速に処理される
  • インデックスタプルの左側がWHEREにくるようにする 参考

kindergartensも同様に、後でMINをとるクエリを使うためにインデックスを張っておきます

ポイント4: トリガーで自動更新

CREATE TRIGGER kindergarten_rate 
  AFTER UPDATE ON infants 
    BEGIN
      UPDATE kindergartens SET rate = 
      (SELECT MAX(rate) FROM infants WHERE kindergarten = old.kindergarten)
      WHERE rowid = old.kindergarten;
      
      UPDATE kindergartens SET rate = 
      (SELECT MAX(rate) FROM infants WHERE kindergarten = new.kindergarten)
      WHERE rowid = new.kindergarten;
      
      INSERT INTO ans 
      SELECT MIN(rate) FROM kindergartens;
    END;

  • infantsテーブルに更新があると自動でkindergartenテーブルとansテーブルに反映されるようにします。こうすることで、いちいちUpdateを発行せずに済み、Python <-> SQLiteのやりとりを減らして高速化できます

  • rowidを用いて高速化します。

    • SQLiteでは行固有のカラムであるrowidが自動で作られています。
    • rowidはinsertのたびに1から自動採番でふられます
    • Insertをうまく行うことで、rowidと幼稚園IDが一致するようにします。
      • これは後述の「ポイント6: PythonループでなくSQLで初期化」で行います

ポイント5: トランザクションで高速化？

cur.execute("begin")
  • トランザクションをはると複数のInsertが高速になるという記事を見かけたので、トランザクションを開始しておきます。
  • commitは不要なのでしません
  • 速くなってるような、なってないようなという感じですがおまじないでいれておきます

ポイント6: stdin.readlineで高速読み取り

    for _ in range(N):
        # a: rate, b: kindergarten
        rate, kindergarten = map(lambda x: int(
            x), sys.stdin.readline().split())
        cur.execute("INSERT INTO infants VALUES(?, ?)", (kindergarten, rate))
  • 標準入力から読み取った値をinfantsテーブルに格納します
  • input()ではなくsys.stdin.readline()を使うと速くなります

ポイント7: PythonループでなくSQLで初期化

WITH RECURSIVE cnt(x) AS (
SELECT 1
UNION ALL
SELECT x+1 FROM cnt
LIMIT 200001
)     
INSERT INTO kindergartens 
SELECT MAX(rate) 
FROM cnt LEFT OUTER JOIN infants
ON cnt.x = infants.kindergarten
GROUP BY x

初期化のポイントは次になります

  • SQLで行うことでPython <-> SQLiteのやりとりを減らして高速化
    • 普通にやるとPythonのループ内でINSERTを何回も発行するが遅い
  • 「ポイント4: トリガーで自動更新」のために、幼稚園IDとDBのrowidが一致するようにする。幼稚園ID順で行が並ぶように、NULLも含めてinsertする

上記を実現するために、次のテクニックを利用します

  • WITH RECURSIVE ...の部分で、1~200001の連番の入っているテーブルcntを作成
  • LEFT OUTER JOIN...することで、1~200001のkindergartenに所属する園児レート(なければNULL)のテーブルを作成します
  • MAX(rate) GROUP BY xで幼稚園IDごとの最大値(なければNULL)をSELECT

あとはクエリを処理して結果出力するだけ

    for _ in range(Q):
        # c: infant, d: next kindergarten
        infant, next_kindergarten = map(
            lambda x: int(x), sys.stdin.readline().split())

        cur.execute("UPDATE infants SET kindergarten = ? WHERE rowid = ?",
                    (next_kindergarten, infant))
    cur.execute("select * from ans")
    ans = cur.fetchall()
    for (result, ) in ans:
        sys.stdout.write(str(result) + "\n")
  • infantsテーブルにUpdateすると「ポイント4: トリガーで自動更新」のトリガーが自動で実行され、kindergartenテーブルとansテーブルに結果が反映されます
  • 全クエリを格納後にansの中身を出力すれば答えです

まとめ

Python標準搭載のSQLiteでSQLを使ってAtCoderの問題が解けることを解説しました。
PythonやSQLの高速化テクニックが必要でめっちゃ大変でした

私のような物好きな人はSQLを使ってみてはどうでしょうか。
追記: 最近もSQLで解ける問題が出たみたいです
https://atcoder.jp/contests/abc217/tasks/abc217_d

追記2:
現在では平方分割によるmultisetライブラリが共有されているので、そっちを使う方が簡単そうです
https://github.com/tatyam-prime/SortedSet

Discussion