配列で渡された要素をグループに分ける

はじめに

class Ball
{
  double Radius;
  RGB Color;
}

みたいなクラスがあって、IEnumerable<Ball> を色別にグルーピングする必要がありました。
グルーピング後に「ある色について何かする」という処理がたくさんあり、扱いやすく効率的なものにしたいと思いました。

IEnumerable<Ball> balls = ...;

// 色別にグルーピング
var grouped = balls.Grouping(b => b.Color);

// 赤色のボールについて何かする
IEnumerable<Ball> redBalls = grouped.GetItem(RGB.Red);
foreach (Ball ball in redBalls) {
  ...
}

// 緑色のボールを削除
grouped.Remove(RGB.Green);

GroupBy()

最初に思いついたのは、LINQ の GroupBy() を使う方法です。
特定のグループを取り出すには、Where() とか Single() が使えます。

IEnumerable<Ball> balls = ...;

// 色別にグルーピング
var grouped = balls.GroupBy(b => b.Color);
// grouped の型は IEnumerable<IGrouping<RGB, Ball>>

// 赤色のボールについて何かする
var redBalls = grouped.Single(g => g.Key == RGB.Red);
// redBalls の型は IGrouping<RGB, Ball> で IEnumerable<Ball> の子クラス
foreach (Ball ball in redBalls) {
  ...
}

// 緑色のボールを削除
grouped = grouped.Where(g => g.Key != RGB.Green);

GroupBy() の返り値の型は IEnumerable<IGrouping<K, T>> になります。
グループの検索は常に先頭から順番にたどっていくことになり、実行時間は $O(n)$ となります。
また削除についても、あるグループを含まないコレクションを作り直すしかなく、ピンポイントで削除できませｎ。
グループに対する操作を何度も繰り返す場合、これはオーバーヘッドが大きすぎます。

ちなみに、.GroupBy().Single() としたときの返り値の型は IGrouping<K, V> です。
これは IEnumerable<V> を継承していますので、あるグループ内の要素についてループを回す分には困りません。

LookUp

キーを使ったグルーピングをする型として LookUp<K, T> があります。
IEnumerable<T>.ToLookUp() で簡単に変換できます。

IEnumerable<Ball> balls = ...;

// 色別にグルーピング
ILookUp<RGB, Ball> grouped = balls.ToLookUp(b => b.Color);
// grouped の型は ILookUp<RGB, Ball>

// 赤色のボールについて何かする
IEnumerable<Ball> redBalls = grouped[RGB.Red];
foreach (Ball ball in redBalls) {
  ...
}

// 緑色のボールを削除
grouped.Remove(RGB.Green); // No such method.

ILookUp<K, T> に変換した後はインデクサーが使えますので、グループの検索は $O(\log n)$ になります（そのような実装になっていると思いたいです）。
また、ILookUp<K, T> は IEnumerable<IGrouping<K, T>> を継承していますので、GroupBy() と同様に扱えます。

ただし、グループを削除する方法がありません。
Where() で削除すると、型が ILookUp<K, V> から IEnumerable<IGrouping<K, V>> に変わってしまいます。

Dictionary

キーを使った検索は辞書( Dictionary<K, V> )型が得意とするところです。
そこで、配列を値とした辞書に格納することにしました。

自分で書いてみる

まずは、自分で Dictionary<RGB, List<Ball>> を組み立ててみました。
ボールの色を順番に調べて、既存の色なら List に加え、未知の色なら辞書にキーを追加して、を繰り返します。

IEnumerable<Ball> balls = ...;

// 色別にグルーピング
Dictionary<RGB, List<Ball>> grouped = new Dictionary<RGB, List<Ball>>();
foreach (Ball ball in balls) {
  List<Ball> list;
  if (!grouped.TryGetValue(ball.Color)) {
    list = new List<Ball>();
    grouped.Add(ball.Color, list);
  }
  list.Add(ball);
}

// 赤色のボールについて何かする
List<Ball> redBalls = grouped[RGB.Red];
foreach (Ball ball in redBalls) {
  ...
}

// 緑色のボールを削除
grouped.Remove(RGB.Green);

できあがってしまえば、グループの検索も削除も $O(\log n)$ になります。
あるグループ内の要素についてループを回すのも簡単です。

でも、自分で Dictionary と List をメンテするのはスマートとはいえません。

ToDictionary()

IEnumerable<T>.ToDictionary() があります。
名前の通り、IEnumerable<T> から Dictionary<K, V> に変換できます。

しかし、IEnumerable<T> の要素1つから Dictionary<K, V> の要素1つに変換するので、このままではグルーピングには使えません。

GroupBy() と ToDictionary()

GroupBy() の戻り値は IEnumerable<> 型なので、ToDictionary() が使えます。
そこで、list.GroupBy().ToDictionary() として辞書にしてみました。

IEnumerable<Ball> balls = ...;

// 色別にグルーピング
Dictionary<RGB, IGrouping<RGB, Ball>> grouped = balls.GroupBy(b => b.Color).ToDictionary(g => g.Key);
// grouped の型は Dictionary<RGB, IGrouping<RGB, Ball>>

// 赤色のボールについて何かする
var redBalls = grouped[RGB.Red];
// redBalls  の型は IGrouping<RGB, Ball>
foreach (Ball ball in redBalls) {
  // IGrouping<RGB, Ball> は IEnumerable<Ball> の子クラス
  ...
}

// 緑色のボールを削除
grouped.Remove(RGB.Green);

辞書型なので、グループの検索も削除も簡単で速い（ $O(\log n)$ ）です。
しかし、グループを取得する際に Dictionary<K, V>.TryGetValue(K, out V) を使おうとすると、C#6.0 以前では型を知っている必要があります。

until6.0.cs

// 黒色のボールがあれば何かする
IGrouping<RGB, Ball> blackBalls;
if (grouped.TryGetValue(RGB.Black, out blackBalls)) {
  foreach (Ball ball in blackBalls) {
    ...
  }
}

after7.0.cs

if (grouped.TryGetValue(RGB.Black, out var blackBalls)) {
  foreach (Ball ball in blackBalls) {
    ...
  }
}

out 引数の時だけグループの型が IGrouping<K, V> であることを思い出すのは面倒くさいです。
いつでも IEnumerable<V> 型で扱いたいので、もう一工夫します。

値をキャストして辞書化

ToDictionary() は、第 2 引数にも関数を渡すことで、要素から値を生成できます。
そこで、値の方を IEnumerable<V> 型にキャストして格納してみます。

IEnumerable<Ball> balls = ...;

// 色別にグルーピング
var grouped = balls.GroupBy(b => b.Color).ToDictionary(g => g.Key, g => g as IEnumerable<Ball>);
// grouped の型は Dictionary<RGB, IEnumerable<Ball>>

// 赤色のボールについて何かする
var redBalls = grouped[RGB.Red];
// redBalls  の型は IEnumerable<Ball>
foreach (Ball ball in redBalls) {
  ...
}

// 緑色のボールを削除
grouped.Remove(RGB.Green);

// 黒色のボールがあれば何かする
IEnumerable<Ball> blackBalls;
if (grouped.TryGetValue(RGB.Black, out blackBalls)) {
  foreach (Ball ball in blackBalls) {
    ...
  }
}

ようやく、扱いやすい型になりました。

おわりに

LINQ と C# 7.0 以降なら、グルーピングはとても簡単で、.GroupBy(e => e.K) とか .GroupBy(e => e.K).ToDictionary(g => g.Key) で済みます。

C# 6.0 以前で辞書で out 引数を使おうとして、ちょっと嵌った話でした。