動的計画法によるDVDのディスク分割の改善
こんにちは。「家族アルバム みてね」の開発チームに所属している黒川と申します。今回は、その「みてね」の機能の1つで、写真や動画をDVDにして注文できる機能を動的計画法を使って改善した話をします。
「みてね」では家族の写真や動画をアップロードし、アプリ上で月ごとに振り返ることが可能になっています。一方、たとえば自宅のテレビやパソコンでまとめて振り返りたいという要望もあり、「みてね」では最長過去1年間の写真や動画をDVDにまとめて注文することができます。
このときに問題となるのがDVDのディスク分割です。1年分の写真や動画はともすると1枚のディスクに収まりきらず、複数のディスクに分割する必要があります。いままでは、動画を月ごとに分けて各ディスクに入れていく、というシンプルなアルゴリズムで分割を行っていました。しかし、ユーザーさんからは「1枚のディスクにすこしの動画しかないがどうなっているのか」「ディスクの枚数をもっと減らせないのか」といった問い合わせが寄せられていました。
そこでディスク分割のアルゴリズムを改善できないかと考えた結果、つぎのように問題を定義しました。
問題設定
N 本の動画の並びが与えられます。また各動画は時間 t （ミリ秒）と、どの年月のものかを表すキー m を持っていて、動画はキーが昇順になるように並べられています。また、ディスク1枚に収まる動画の時間容量 L（ミリ秒）とディスクの総枚数の上限 M が与えられます。
黒川君は動画の順番を保ったまま、必要最小限なディスク枚数に動画を分割して、各ディスクに割り当てていきたいです。ただし各ディスクに含まれる動画の合計時間は L を超えてはいけません。またディスクの総枚数は M 以下でなければなりません。
さらに黒川君はディスク枚数が変わらない限り、できるだけ m が等しい動画は同じディスクに入っていてほしいと考えています。すなわち、各 m についてキーが m であるような動画が含まれるディスクの枚数 - 1を cost[m] とすると、Σcost を最小にしたいです。
N, M, L, t, mが与えられたときに黒川君のかわりに最適な分割を求めなさい。解が複数ある場合はどれを返してもよいものとします。
1 ≦ N ≦ 10⁵
1 ≦ M ≦ 50
1 ≦ L ≦ 5×10⁶
1 ≦ t[i] ≦ L for i = 1, 2, …, N
1 ≦ m[i] ≦ N for i = 1, 2, …, N
メモリ制約：256MB、時間制約：10秒
たとえば N = 3, M = 50, L = 50, t = {40, 10, 30}, m = {1, 2, 2} であれば、{40, 10}, {30} と分割すると cost = {0, 1} となり cost の総和は 1 になりますが、 {40}, {10, 30} と分割すれば cost = {0, 0} となり、cost の総和は 0 でこちらの方が小さくなります。どう分割してもディスクはすくなくとも2枚必要なので、後者の分割が解となります。
あるいは、N = 4, M = 50, L = 50, t = {30, 10, 30, 20}, m = {1, 2, 2, 3} の場合、分割は {30, 10}, {30, 20} の分割が最適となり cost の総和は 1 となります。ディスク枚数を4枚にして {30}, {10, 30}, {20} とすれば cost の総和は 0 になりますが、ディスク枚数は必要最小限にする、という条件があるので、これは解にはなりません。
解法（愚直バージョン）
さて、この問題を解いてみましょう。まず思いつくのは次のような方法です：まず何も考えずに動画をディスクに順に詰めていきます。もしディスクが一杯になったら新しいディスクに詰めます。すると最終的に得られるディスクが必要最低限なディスクの枚数になるはずです。これを M' としましょう。
すると、あとは動画のどこをディスクの区切りにすれば良いかを選べばよいことになります。たとえば動画 1 と動画 2 の間、動画 3 と 動画 4 の間を区切ることにすれば分割は {1}, {2, 3}, {4, …}となるでしょう。それぞれの場合について cost の総和を計算して一番小さくなった分割を選べばよいことになります。
この方法の計算量を見積もってみましょう。区切りの選び方の総数は、二項係数を C(i, j) と書くことにすると C(N-1, M'-1) になり大雑把に N! になってしまいます。Nが12でも ~10⁸ くらいになってしまい、とても数秒では終わらなくなってしまいます。なにかよい方法を考える必要があります。
解法 （動的計画法バージョン）
ここで、動的計画法が使えないか検討してみます。動的計画法というのはいま考えている問題の部分問題の結果を記憶しておくことで全体の計算を節約する手法のことです。
ここで、つぎのような関数を定義します。
dp(i, j) := i枚のディスクに1~j番目までの動画を詰めたときの最小コスト
ただし、i枚のディスクに1~j番目の動画を詰められないときは∞になるとします。また、0枚のディスクには（便宜上）0番目までの動画を詰めることができると考え、dp(0, 0) = 0とし、それ以外の i, j については ∞ を初期値を設定します。すると漸化式はつぎのようになります。
漸化式の内容は簡単にいえばつぎのようになります。まず dp(i, j) を計算するときに1~jのどこかが最後の区切りになるはずです。その箇所をkとします。すると、k~jの後半部分の動画をi枚目のディスクに入れ、1~k-1の前半部分の動画をi-1枚のディスクに入れることになります。すると前半部分のコストの最小値は dp(i -1, k-1)になります。一方、後半部分のコストは、もしk~jが1枚のディスクに収まるなら∞に、そうでなければm[k-1]とm[k]が異なる場合のみ 1 加算されます。両者を足し合わせれば dp(i, j) が求まります。
最終的に dp(i, N) のうち∞にならない最小の i を求め、あとはそのコストが達成できる分割を逆にたどっていけば解が得られます。
iとjを小さい値から順に計算していくとすると、漸化式で現れる dp(i-1, k-1) はすでに計算済みなので計算時間はかかりません。またk~j番目の動画が1枚のディスクに収まるかどうかはあらかじめ t の累積和を計算しておけば定数時間で計算できます。各漸化式の計算では k を 1 から j まで動かすので計算量は O(N)、さらに i と j を動かすので、全体では O(MN²) かかることになります。
ここで N のオーダーが10⁵であったことを思い出すと全体で10¹¹になってしまい、とても数秒では終わりません。まだ何らかの改良が必要です。
解法（動的計画法 + しゃくとり法 + 優先度付きキューバージョン）
ここで、さきほどの解法で無駄な部分がないか考えてみましょう。k を 1 から j まで動かして最小値を探している箇所はどうでしょうか。k の探索範囲を減らすことができれば全体で O(MN) 程度にできるかもしれません。
まず、第2項について考えましょう。k ~ j の動画が1枚のディスクに収まるかチェックしていますが、ここはしゃくとり法で簡単にできそうです。しゃくとり法というのは、競技プログラミングでよく出てくるテクニックで、何らかの条件を満たす区間を数え上げるのに使うものです。詳細はQiitaの解説記事に譲りますが、今回の例では、k ~ j の動画が1枚のディスクに収まるような区間を考えたときに、その左端、つまり k の取り得る範囲の最小値を j の関数として f(j) とすると、f(j) は j に対して広義単調増加になっている点がポイントになっています。
しゃくとり法は簡単にいうと、j の前回の結果をうまく利用することで k の探索範囲を狭めるものです。具体例で見てみましょう。
いま L = 50, t = {10, 20, 20, 20, 30}, j = 4 とします。まず k = 1のとき、つまり一番上の図の状態では k ~ j の動画の時間の合計は 70 で L に収まりません。そこで k ~ j の合計が 50 以下になるまで k を進めていきます。すると k = 3 のところで合計が 40 になります（2番目の図）。つまり j = 4 のときは k の探索範囲の下限は 3 ということになります。その状態で漸化式の第2項を計算し dp(i, j) を計算します。つぎに j を 1 増やします。このとき、さきほどのように k = 1 から探索を開始するのではなく、さきほど探索を終えた k = 3 から調べます（3番目の図）。なぜなら、j = 4 のときの k の探索範囲の下限が 3 だったのですから、 j = 5 のときの k の探索範囲の下限はすくなくとも 3 以上だからです。このようにして k ~ j が1枚のディスクに収まるかどうかは効率的にチェックできるようになりました。
しかし、まだ続きがあります。漸化式の最小値を見つけるためには依然として k を動かさなければなりません。これについては優先度付きキューを用います。優先度付きキューは、キューの要素の中から最小値もしくは最大値をもったものを O(log N) で得ることができるデータ構造です。これを使って各 i において j を動かしつつ同時に dp(i -1, j - 1) を優先度付きキューに入れておけば、漸化式の計算をするタイミングでそのキューから取り出した最小値が求める第1項になります。あとはさきほどのしゃくとり法と組み合わせて k の探索範囲の下限を毎回更新してあげるときに、キューから用済みになった要素を取り除いてあげれば、全体として O(MN log N) で計算できます。
以上のアルゴリズムをコードで書くとこんな感じになります。Rubyなので変数名に大文字が使えないため、ちょっと分かりづらいですが、雰囲気が伝われば幸いです。
さて、以上で計算量は O(MN log N) になり、オーダーとしては 10⁶ 程度なので Ruby でも数秒程度で計算できるはずでした。しかし、実際に手元で試したところ、予想される最大サイズのデータで試すと20秒以上かかってしまいました。どうやら PriorityQueue の呼び出しに時間がかかっていたようです。なお、自分は普段はJava/Kotlinを書いているので、最初は慣れたJavaで実装したのですが、Java実装では数秒程度で実行できたので、この結果はいささか意外でした。
この記事では紙面の関係で触れませんが、その後、優先度付きキューを使わずに代わりのデータ構造を使うことで全体の計算量を O(NM²) にすることができました。結果として、計算は数秒程度で終わり、いまもプロダクション環境で元気に動いています。また、ユーザーさんからの問い合わせもいままで月に数回程度あったものが、アルゴリズム改良後は3ヶ月で1つもないので、ある程度ユーザーさんの満足度に貢献できたのではないかと思います。
振り返り
自分は1年ほど前から趣味で競技プログラミングをしているのですが、今回、業務でその知識を活かせそうな機会があったので、マネージャーにお願いしてやらせてもらいました。
実際にやってみて感じたこととしては、「そもそも問題の形に落としこむまでが大変」ということです。この記事では定式化された問題を天下り的に示しましたが、実際は自分がそのドメインに詳しくなかったこともあって、前提や情報が不十分なままアルゴリズムを考え始めて何度か失敗してしまいました。競技プログラミングをすこしでもやったことがある方は分かると思いますが、最適なアルゴリズムを考える上で入力の制約がすこし違っただけでもまったく異なる解になることがしばしばです。また、そもそもどのような分割が望ましいか、という数字にしにくい点についてプロダクトオーナーと合意を得て数式の形に落とすことも重要なポイントだったと思います。
もう1つの教訓としては、使用する言語やフレームワークが重要という点です。笑い話に聞こえると思いますが、Javaで実装して数秒で実行できたときは、あとはRailsに移すだけだから1日で終わると思っていました。実際は前節で述べたように、さらなるアルゴリズムの改良も必要でしたし、また不用意にActiveSupportのメソッドを呼んだせいでびっくりするくらい遅くなったり、テスト実行時にTime.zone.parseを呼んだらとても遅くなったり、とさまざまなことをやらかしました。結果として、Railsに移植するだけで1週間くらいかかっています。こういったパフォーマンスに厳しいコードを書くときはかなり気を遣う必要があることを学びました。なお、ボトルネックの調査にはstackprofというgemがほぼオーバーヘッドなしでどのメソッドが遅いか調査できて大変便利でした。
そのむかし The Algorithm Design Manual という本を夢中になって読んでいたときにたまに挟まる War Story というコラムがとても好きでした。これは著者の Steven S. Skiena さんが、実際の業務で遭遇したアルゴリズムの問題を解決する話なのですが、たいへんエキサイティングでおもしろく、いつかは自分もそういう仕事ができたらと思っていたので、今回ちょっと似たような経験ができて満足でした。この場を借りて、そのような機会を与えてくれたマネージャーやサポートしてくれた同僚、そしていつも「みてね」を使ってくれるユーザーさんに感謝いたします。
