Dictの値をキーにキーを集計し、集計されたキーに対して任意の関数を実行したDictを返す

Elixir Advent Calendar 2014 12日目。
お仕事が多忙を極めている(一年中この状態という説もある)上に、忘年会が沢山あってつらみが増しているので、今日はライトなネタを。

やりたいこと

あるHashed Listがあるとします。

list = [a: 1, b: 0, c: 4, d: 0, e: 1, f: 1]

これに関して、バリューをキーにして、集計したい。
期待するアウトプットは、こんな感じです。

%{0 => ["b", "d"], 1 => ["a", "e", "f"], 4 => ["c"]}

ちょっとしたデータ処理に便利ですよね。
標準関数にはちょうどいいのが無いので、作ってみました。

Enum.group_by/3 の拡張

ちなみに、近い関数としては Enum.group_by/3 があるのですが、それだけでは上記は実現できません。もう一つ処理が必要です。

まず、 Enum.group_by/3 のソースを見てみましょう。

@spec group_by(t, dict, (element -> any)) :: dict when dict: Dict.t
def group_by(collection, dict \\ %{}, fun) do
  reduce(collection, dict, fn(entry, categories) ->
    Dict.update(categories, fun.(entry), [entry], &[entry|&1])
  end)
end

ちょっと足りませんね。
これでどうでしょう。

@spec group_by(t, dict, (element -> any), (element -> any)) :: dict when dict: Dict.t
def group_by(collection, dict \\ %{}, fun, fun_entry) do
  Enum.reduce(collection, dict, fn(entry, categories) ->
    Dict.update(categories, fun.(entry), [fun_entry.(entry)], &[fun_entry.(entry)|&1])
  end)
end

第四引数に、関数をもう一つ追加しました。

実行結果

defmodule EnumX do
  def group_by(collection, dict \\ %{}, fun, fun_entry) do
    Enum.reduce(collection, dict, fn(entry, categories) ->
      Dict.update(categories, fun.(entry), [fun_entry.(entry)], &[fun_entry.(entry)|&1])
    end)
  end
end

EnumX.group_by([a: 1, b: 0, c: 4, d: 0, e: 1, f: 1], fn({_, v}) -> v end, fn({k, _}) -> to_string(k) end)

%{0 => ["d", "b"], 1 => ["f", "e", "a"], 4 => ["c"]}

うまくできました。

応用

この関数、わりと汎用性が高いと思っています。このような使い方もできるでしょう。

EnumX.group_by(~w{ant buffalo cat dingo}, &String.length/1, &String.capitalize/1)

%{3 => ["Cat", "Ant"], 5 => ["Dingo"], 7 => ["Buffalo"]}

第三引数までは公式ドキュメント通りですが、与える関数を追加することにより、それぞれをキャピタライズしているところが違います。
標準関数であっても、パイプ演算子を使うことで実現できるとは思いますが、力技になるのではないでしょうか。（この例の場合、第一引数を Enum.group_by/3 に与える前に String.capitalize/1 しとけばいい話ですけど。）

まとめ

Dictの値をキーにキーを集計し、集計されたキーに対して任意の関数を実行したDictを返す方法をご紹介しました。

明日は @ma2ge さんです。