【C#】セグメント木を実装してみる - はなちるのマイノート

はじめに

今回はセグメント木を実装してみようという記事になります！

セグメント木は主に区間上の値の更新と任意の区間内の最小値などの取得を高速化できます。

構造としては完全二分木を用いていて、初期化に $o(n)$ ，更新・取得は共に $o(logn)$ で動作します。
en.wikipedia.org

さっそくみていきましょう。

はじめに
仕組み
コード
使い方
実例
さいごに

仕組み

この記事で細かいところまで紹介をしたかったのですが、以下の記事がとても分かりやすかったのでそちらを参照してみてください。

tsutaj.hatenablog.com

プログラミングコンテストでのデータ構造

基本方針はこれらの記事の実装を参考にしますが、より柔軟に使えることを目指してコードを書いてみました。

コード

public class SegmentTree<T>
{
    public int N { get; private set; }
    private T[] _tree;
    private readonly Func<T, T, T> _updateMethod;
    private readonly T _initValue;

    /// <summary>
    /// セグメント木の初期化
    /// </summary>
    /// <param name="items">要素達</param>
    /// <param name="updateMethod">更新する方法</param>
    /// <param name="initValue">セグメント木のノードの初期値</param>
    public SegmentTree(IEnumerable<T> items, Func<T, T, T> updateMethod, T initValue) {
        T[] array = items.ToArray();
        _updateMethod = updateMethod;
        _initValue = initValue;

        // セグ木の最下段は２の冪乗にする
        N = 1;
        while (N < array.Length) N *= 2;

        _tree = Enumerable.Repeat(initValue, 2 * N - 1).ToArray();

        // 最下段に要素達を入れた後、下の段から更新していく
        for (int i = 0; i < array.Length; i++) _tree[N + i - 1] = array[i];
        for (int i = N - 2; i >= 0; i--) _tree[i] = _updateMethod(_tree[2 * i + 1], _tree[2 * i + 2]);
    }

    /// <summary>
    /// 更新する
    /// </summary>
    /// <param name="i">更新したい値のインデックス</param>
    /// <param name="x">更新する値</param>
    public void Update(int i, T x)
    {
        // 更新したい要素のセグ木上でのインデックスを取得
        i = N + i - 1;

        // 値を更新した後に、どんどん親を更新していく
        _tree[i] = x;
        while(i > 0)
        {
            i = (i - 1) / 2;
            _tree[i] = _updateMethod(_tree[2 * i + 1], _tree[2 * i + 2]);
        }
    }

    /// <summary>
    /// 区間内の目的の値を取得する，要求区間[a, b)
    /// </summary>
    /// <param name="a">a以上の区間</param>
    /// <param name="b">b未満の区間</param>
    /// <returns></returns>
    public T Execute(int a, int b)
        => Execute(a, b, 0, 0, N);

    private T Execute(int a, int b, int k, int l, int r)
    {
        // 要求区間[a, b)に対して対象区間[l, r)を求める
        // 今いるノードのインデックスがk

        // 要求区間と対象区間が交わらない
        if (r <= a || b <= l) return _initValue;

        // 要求区間が対象区間を完全に被覆
        if (a <= l && r <= b) return _tree[k];

        // 要求区間が対象区間の一部を被覆しているので、子について探索を行う
        // 新しい対象区間は、現在の対象区間を半分に割ったもの
        var vl = Execute(a, b, 2 * k + 1, l, (l + r) / 2);
        var vr = Execute(a, b, 2 * k + 2, (l + r) / 2, r);
        return _updateMethod(vl, vr);
    }
}

使い方

区間内の最小値を求めたい場合はこんな感じ。

static void Main(string[] args)
{
    // {1, 2, 3, 4, 5}からセグメント木を作成し、更新方法を最小値に設定する
    // initValue(一番右の奴)はセグ木のノードの初期値のことで、最小値ならとりあえず大きい値を入れればOK
    var tree = new SegmentTree<int>(Enumerable.Range(1, 5), (x, y) => Math.Min(x, y), int.MaxValue);

    // 0以上6未満での最小値を求める
    Console.WriteLine(tree.Execute(0, 6));  // 1

    // 1以上4未満での最小値を求める
    Console.WriteLine(tree.Execute(1, 4));  // 2

    // 0番目の値を6に変更する
    tree.Update(0, 6);

    // 0以上6未満での最小値を求める
    Console.WriteLine(tree.Execute(0, 6));  // 2
}

また小数かつ区間の合計を求めるなんかもできます。

static void Main(string[] args)
{
    var nums = new float[] { 1.0f, 2.0f, 3.0f, 4.0f, 5.0f };

    // {1, 2, 3, 4, 5}からセグメント木を作成し、更新方法を合計に設定する
    var tree = new SegmentTree<float>(nums, (x, y) => x + y, 0);

    // 0以上6未満での合計を求める
    Console.WriteLine(tree.Execute(0, 6));  // 15

    // 1以上4未満での合計を求める
    Console.WriteLine(tree.Execute(1, 4));  // 9

    // 0番目の値を6に変更する
    tree.Update(0, 6);

    // 0以上6未満での合計を求める
    Console.WriteLine(tree.Execute(0, 6));  // 20
}

実例

こちらのサイトにてより高度なテストができたりします。
http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=DSL_2_A&lang=jp

static void Main(string[] args)
{
    var nq = Console.ReadLine().Split(' ').Select(int.Parse).ToArray();
    var n = nq[0];
    var q = nq[1];
    var tree = new SegmentTree<int>(Enumerable.Repeat(int.MaxValue, n), (x, y) => Math.Min(x, y), int.MaxValue);
    var results = new List<int>();
    for (int i = 0; i < q; i++)
    {
        var inputs = Console.ReadLine().Split(' ').Select(int.Parse).ToArray();

        if (inputs[0] == 0)
        {
            tree.Update(inputs[1], inputs[2]);
        }
        else if (inputs[0] == 1)
        {
            var result = tree.Execute(inputs[1], inputs[2] + 1);
            results.Add(result);
        }
    }

    foreach (var item in results) Console.WriteLine(item);
}