yield return

C#

イテレータを簡単に書けるyield returnが便利で好きなのだけれど、以下のようなのが書けたらなーと思ったのでメモ。

using System;
using System.Collections.Generic;

class Hoge : IEnumerable<int>
{
  int[] _a;
  int _b;
  int[] _c;

  public Hoge(int[] a, int b, int[] c)
  {
    _a = a;
    _b = b;
    _c = c;
  }

  public IEnumerator<int> GetEnumerator()
  {
    if (_a != null) return _a.GetEnumerator();  // returnでもyield returnでもエラー
    yield return _b;
    if (_c != null) return _c.GetEnumerator();  // 同上
  }

  System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
  {
    return GetEnumerator();
  }
}

class App
{
  public static void Main(string[] args)
  {
    int[] a = new int[]{1, 2};
    int b = 3;
    int[] c = new int[]{4, 5, 6};
    Hoge h = new Hoge(a, b, c);

    // 1 2 3 4 5 6 って走査できたらなぁ...
    foreach(int n in h)
    {
      Console.Write(n + " ");
    }
  }
}

これができると、木構造のトラバースが簡単になるなー、みたいな。こんな感じ。

// 2分木のノード
class Node : IEnumerable<int>
{
  int _val;
  Node _left;
  Node _right;

  // ... 略 ...

  public IEnumerator<int> GetEnumerator()
  {
    if (_left != null) return _left.GetEnumerator();
    yield return _value;
    if (_right != null) return _right.GetEnumerator();
  }
}

// 2分木
class Tree : IEnumerable<int>
{
  Node _root;

  // ... 略 ...

  public IEnumerator<int> GetEnumerator()
  {
    if (_root == null) yield break;
    else
    {
      yield return _root.GetEnumerator();
    }
  }
}

要素をなめるだけなら

class Node
{
  // ... 略 ...

  void ForEach(Action<int> action)
  {
    if (_left != null) _left.ForEach(action);
    action(_val);
    if (_right != null) _right.ForEach(action);
  }
}

とかやれば良いけど、IEnumerableにしたい。TreeNodeEnumeratorクラスを作るのを面倒くさく思ったのがキッカケだけど、気付いてないだけで簡単な方法がありそうだな...。うーん...。

うほ！いい記事！