本記事は Wantedly 21新卒 Advent Calendar の9日目の記事です。本記事では、競技プログラミングで良く用いられるデータ構造、素集合データ構造の OCaml による実装を行います。またその実装を通じて、適切な抽象化機構の備わったモダンな言語で競技プログラミングに取り組む際の、僕のアプローチを紹介します。

　競技プログラミングが人口に膾炙するに伴ってか僕が社会人になってから精進を怠っていたからか、最近では競プロ勢のレベル向上に驚くことばかりです。以前ならこの問題が解ければ水色くらいのパフォーマンスは出ただろうと思しき所で緑色程度しか出なかったりと、様々なデータ構造やアルゴリズムの知識が常識として共有されてきているのを感じます。

素集合データ構造、通称 union-find も競プロ勢に広く知られたデータ構造の一つでしょう。

素集合データ構造とは

　素集合データ構造とは互いに素な集合を管理するデータ構造で、2つの要素が与えられた際にそれらが同じ集合に属するかを求めるクエリ、及び2つの集合を合併して1つの集合にする操作を高速に処理できます。典型的には、要素を頂点とし、同じ集合に属する要素を辺で繋いだ根付き木の集合として素集合データ構造を実装することが多いです。

　根付き木の集合として素集合データ構造を表現する場合、まず2つの要素が与えられた際にそれらが同じ集合に属するかを求めるクエリは、2つの要素に対応する頂点から親を辿って、根同士が等しいがどうかで判別できます。ここで親を辿って根を見つけた際、それまでに通った頂点の親を根に付け替えてショートカットしてしまうことで、次回以降のクエリを高速化するのが一般的です。

　また、2つの集合を合併して1つの集合にする際は、集合に対応する木の根を求め、一方の根からもう一方の根へ辺を貼ることで実装できます。この時、高さが低い方の木の根を親にすることで、前述のクエリの時間計算量がより改善します。

　あまり self-contained な記事を目指すと執筆が大変なので非常にざっくりとした説明になりましたが、より詳しく知りたい場合は AtCoder 社による解説 であったり、プログラミングコンテストチャレンジブック、通称蟻本であったりを読むのが良いでしょう。

素集合データ構造の実装

　素集合データ構造がどういったものか解説したところで、実際に OCaml で実装していきましょう。

　AtCoder 社による解説 や プログラミングコンテストチャレンジブック でも採用されているように、根付き木の集合として素集合データ構造を表現する際には、集合の要素に対応する添字に、親にあたる頂点の添字を格納した配列を用いるのが一般的です。根にあたる頂点は親を持たないので、番兵を入れておくと良いでしょう。

　まず初期化処理、n 個の単集合からなる素集合データ構造を作成する処理は、単純に番兵のみが入った n 要素の配列を作れば良いです。配列の添字は 0 から n-1 の整数値を取りますから、負の数か n 以上の正整数であれば番兵としての用をなすことでしょう。折角なので符号を反転した木の高さを番兵として採用することで、高速化のための情報を残しておきます。

(* n 個の単集合からなる素集合データ構造を作成する *)
let init : int -> int array = fun n ->
  (* 単一の頂点からなる木の高さは 1 なので、番兵として入れる値は -1 *)
  Array.make n (-1)

　次に、2つの要素が同じ集合に属するかの判定クエリですが、これは要素に対応する頂点から親を辿って、根にあたる頂点に対応する添字を求めるヘルパ関数、root を用意すれば容易に実装できます。ここで root が根を見つけた際に、それまで通った頂点の親を根に付け替えることで高速化できます。

(* 配列で実装された素集合データ構造とその要素に対応する添字が与えられた時に、
   親を辿って根にあたる頂点の添字を返す関数 *)
let rec root : int array -> int -> int = fun a i ->
  if a.(i) < 0
  then i
  else
    (* 根に対応する添字は j *)
    let j = root a i in
    (* この呼び出しで通った頂点の直接の親を根に付け替える *)
    a.(i) <- j;
    (* 根に対応する添字を返す *)
    j

(* 配列で実装された素集合データ構造と、2つの要素に対応する添字が与えられた際に、
   その2要素が同じ集合に属するかどうかを判定する *)
let equal : int array -> int -> int -> bool = fun a i j ->
　　 (* 根に対応する添字が等しいかどうかで判定できる *)
  root a i = root a j

　最後に、2つの集合を合併して1つの集合にする操作ですが、これもまたヘルパ関数 root を用いて比較的容易に実装できます。番兵の値に木の高さを埋め込んだことですし、その情報も使って高速化もできますね。

(* 配列で実装された素集合データ構造と、2つの合併したい集合の要素に対応する添字が与えられた際に、
   その2要素が属する集合同士で合併して1つの集合にする *)
let unite : int array -> int -> int -> unit = fun a i j ->
  let k = root a i in
  let l = root a j in
  (* 元から同じ集合だった場合は何もしなくて良い *)
  if k <> l then
    (* 添字 k　に対応する頂点が根になるような木の高さ *)
    let hk = ~- (a.(k)) in
    (* 添字 l　に対応する頂点が根になるような木の高さ *)
    let hl = ~- (a.(l)) in
    if hk < hl
    (* 添字 k　に対応する頂点が根になるような根付き木の方が低い場合は、
       そちらから l に対応する頂点へ辺を貼る *)
    then a.(k) <- l
    (* 添字 l　に対応する頂点が根になるような根付き木の方が低い場合 *)
    (* 元々の高さが同じだった場合、 l に対応する頂点から k に対応する頂点へ伸ばした
       辺の分、高さが1増える *)    
    else (a.(l) <- k; a.(k) <- ~- (hk + if hk = hl then 1 else 0))

換骨奪胎

　さて、ここまで一般的な素集合データ構造の実装を踏襲してきた訳ですが、OCaml に慣れた人が上のコードを読むと思ったことでしょう。OCaml らしい実装ではないな、と。

　まず、根付き木の集合を配列で表現するのが OCaml らしくないです。C++ ばかり使っている競技プログラマーならいざ知らず、我々は GC のある言語を使っているのだから貧乏くさいことはせず、ポインタで辺を表現すべきです。頂点を参照セルで、子から親への辺をポインタで表現して素集合データ構造を実装するならば、以下の通りでしょうか。

(* 素集合データ構造を根付き木の集合として表現する際の、頂点の表現 *)
type t = node ref (* 親にあたる頂点を格納する参照 *)
and node =
  | Root of int (* その頂点が根だった場合は親の代わりに根付き木の高さを入れておく *)
  | Link of t (* 親にあたる頂点 *)

(* 単集合に属する要素に対応する頂点を作る関数 *)
let init () = ref (Root 1)

(* 与えられた頂点から親を辿って根を返す関数 *)
let rec root : t -> t = fun node ->
  match !node with
  | Root _ -> node
  | Link parent ->
      let root = root parent in
      node := Link root;
      root

(* 2つの頂点に対応する要素が同じ集合に属しているか判定する関数 *)
let equal : t -> t -> bool = fun node node' ->
  root node == root node'

(* 2つの頂点に対応する要素が、それぞれ属している集合を合併する関数 *)
let unite : t -> t -> unit = fun node node' ->
  (* node が属する木の根 *)
  let r = root node in
  (* node' が属する木の根 *)  
  let r' = root node' in
  (* r を根とする木の高さ *)
  let Root hr = !r in
  (* r' を根とする木の高さ *)
  let Root hr' = !r' in
  if r != r' then
    if hr < hr'
    then r := Link r'
    else (r' := Link r'; r := Root (hr + if hr = hr' then 1 else 0))

　このように配列ではなく参照セルの集合として素集合データ構造を実装することで、最初に全体集合の要素数を宣言しなくても適宜頂点を追加できるメリットがあります。

let rednodes = Array.init 5 @@ fun _ -> init ()
let blacknode = init () (* 素集合データ構造に後から新しく頂点を追加できる *)
(* 後から追加した頂点に対しても操作を行える *)
unite rednode.(2) blacknode

　加えて、集合の要素から素集合データ構造で対応する頂点を検索するデータ構造に連想配列等を使えば、文字列を要素として持つような素集合データ構造なんかも作れるでしょう。配列による実装が 0 から n-1 の整数しか扱えなかったのとは大違いです。

抽象化

　また、頂点を参照セルで表現する実装を採用してもなお、このコードの再利用性が低いのが気になります。素集合データ構造を使うコードではしばしば、合併後の集合の要素数であったり、代表元であったりが必要になることもありますが、そのような場合にいちいち実装を変更しなければならないのでしょうか？

　OCaml にはファンクターに代表される抽象化機構がありますから、そのように合併後の集合について様々な情報が欲しいといった需要にも対応できます。ファンクターで実装を抽象化した集合を持っておいて、それを返せば良いだけです。

module DisjointSetUnion
  (* 集合の実装 *)
  (Set : sig
    type t
    (* 互いに素な集合同士の和集合を返す関数 *)
    val union : t -> t -> t
  end)
= struct
  type t = node ref
  and node =
    (* その頂点が根だった場合は、親の代わりに根付き木の高さとその集合のデータを入れておく *)
    | Root of int * Set.t
    | Link of t

  (* 与えられた集合に対応する頂点を作る関数 *)
  let init : Set.t -> t = fun s -> ref (Root (1, s))

  (* 抽象化前と実装が変わらないので中略 *)

  let unite : t -> t -> unit = fun node node' ->
    let r = root node in
    let r' = root node' in
    (* s は、r を根とする木が表現している集合 *)
    let Root (hr, s) = !r in
    (* s' は、r を根とする木が表現している集合 *)
    let Root (hr', s') = !r' in
    if r != r' then
      if hr < hr'
      (* 素集合を合併したので、根付き木が表現している集合の情報も更新しておく *)
      then (r := Link r'; r' := Root (hr', Set.union s s'))
      else (r' := Link r'; r := Root (hr + (if hr = hr' then 1 else 0), Set.union s s'))

  (* 与えられた頂点がどんな集合に属しているかを返す関数 *)
  let of_set : t -> Set.t = fun node ->
    let Root (_, s) = !(root node) in
    s
end

　集合を合併する度に和集合を取る操作が行われているのでパフォーマンス面が気になるかもしれませんが、実際に必要となる情報だけ引き回すような実装を渡してやれば大丈夫です。いわゆる shallow embedding というやつですね。

(* 集合の要素数だけ欲しい場合の使用例
   DSUCard.of_set は集合の要素数を返すようになる *)
module DSUCard = DisjointSetUnion (struct
  (* 集合の要素数の情報だけ引き回す *) 
  type t = int
  (* 互いに素な集合2つの要素数が与えられた時に、それら和集合の要素数は単純に和を取ったものになる *)
  let union = ( + )
end)

(* 文字列型の集合の代表元が欲しい場合の使用例
   DSURepr.of_set は集合の代表元を返すようになる *)
module DSURepr = DisjointSetUnion (struct
  (* 集合の代表元の情報だけ引き回す *)
  type t = string
  (* 互いに素な集合2つの代表元が与えられた時に、それら和集合の代表元が欲しければどちらかを返せば良い *)
  let union x _ = x
end)

まとめ

　本記事では、競技プログラミングで良く用いられるデータ構造、素集合データ構造の OCaml による実装を行いました。初めはいかにも C++ 等で書かれたコードの移植といった風情の実装でしたが、改良を経て OCaml らしい、汎用的かつ使いやすいインターフェースを手に入れたのではないでしょうか。

　仕事で様々な言語を書いていても思うのですが、やはりそれぞれの言語にはその言語に合った書き方というものがあります。C++ 以外の言語で競技プログラミングに取り組んでいる人は、一度立ち止まってその言語らしい解法とは何か考えてみても良いのかもしれないですね。