セグメント木で時系列データの統計量を計算する

何気なく @chokudai さんのタイムラインを見ていたら「セグ木」という言葉が目にとまった。

セグ木だけやたらみんな解けるようになってる話、ちょっと考えれば、「もっと高度なデータ構造であるmapを使った問題」がもっと解かれているのだから、標準ライブラリに近しい位置になっちゃえば解かれる、というのはまぁ妥当ではある気がする。
（が、おじさんからするとびっくり）

— chokudai(高橋 直大)@AtCoder (@chokudai) 2026年1月24日

聞いたことのない言葉だったので調べてみると。

セグ木とは「セグメント木」の略称で、配列の任意の範囲の集計値を効率的に計算できるデータ構造

とのこと。
じゃあ時系列データの統計量の集計に使えるじゃん？と思ったのでやってみる。

ちなみに今回作ったデモプログラムがここにある。

github.com

目次

セグメント木とは

文字通り木構造のデータ構造。配列からセグメント木を構築して保持しておくことで、元の配列の任意範囲に対しての集計値の計算が効率的に実行できる。

例えば [1, 2, 3, 4, 5, 6, 7, 8] という配列が与えられたとき。配列の値を終端に持つ二分木を考える。

二分木の各ノードには、それよりも下の値に対して集計したい値を事前計算して持たせておく。
上の図では左右のノードの合計値を持たせている。

見てわかる通り、木の root には [1, 2, 3, 4, 5, 6, 7, 8] の合計値である sum: 36 が計算済みで持たせてある。なので「全範囲の合計は？」という問いに対して root に持たせた値を参照するだけで瞬時に答えらる。

root のすぐ下には「配列の前半の合計をもつノード」と「配列の後半の合計を持つノード」がある。なのでノードを一つたどるだけで「前半の合計」や「後半の合計」についても知ることができる。

では [1, 2, 3, 4, 5, 6, 7, 8] の中央部分の「3~6 の合計」を問われたらどうだろう。

「3~6 の合計」を直接保持しているノードはない。しかし「前半の後半」と「後半の前半」を合算したものが「3~6 の合計」になるので、高々2つのノードを参照して合計を計算すれば「3~6 の合計」が求まることになる。

このように、配列を愚直に走査して計算をすると O(N) の計算時間がかかるところを、事前計算したノードをうまくピックアップして計算することで O(log N) の時間に抑えられる。便利だ。

合計以外の集計にも使える

上記の例では合計値を例にしたが、合計以外の集計にも使える。集計したい値ごとにノードに持たせる値を変えればいい。

• ノード: max(left, right) → 任意範囲の最大値集計
• ノード: cond(left) || cond(right) → 任意範囲に条件 cond() を満たす要素があるか調べる
• ノード: gcd(left, right) → 任意範囲の最大公約数 (Greatest Common Divisor) を集計

そして工夫次第では「平均値」や「分散」のように単純な合算で計算できない数値にも使える。たとえば分散 s² の定義は、

• 

で、s² 自体は単純な合算では求まらない。(「分散の平均」は全体の分散にはならない)

一方で

• 要素数: n
• 総和:
• 二乗総和:

は合算可能なので、これをノードに持たせる。各値から分散を求めるのは最後の最後にやればいい。

並び順に意味のあるデータに最適

セグメント木は start~end で範囲を指定できるような、要素の並び順に意味のあるデータに対して有効だ。
逆に順不同な集合 (Set) に対してランダムアクセス的に集計するのには向いていない。

今回は並び順に意味のあるデータとして時系列データを扱ってみる。

サンプルデータ

適当な時系列データがほしいので Kaggle の Hourly Energy Consumption データセットを使ってみる。

www.kaggle.com

アメリカの電力消費量を1時間ごとに記録した CSV がダウンロードできる。
16年分、約12万行のデータがある。

	Datetime	AEP_MW
1	2004-12-31 01:00:00	13478.0
2	2004-12-31 02:00:00	12865.0
3	2004-12-31 03:00:00	12577.0
4	2004-12-31 04:00:00	12517.0
...	...	...

このデータからセグメント木を構築して平均・分散を計算してみることにする。

セグメント木を実装する

二分木

今回は Rust で実装する。まずは素朴な二分木を表現する構造体を作ろう。

struct Node {
    range: Range<NaiveDateTime>,  // このNodeがカバーする日時範囲
    stats: Stats,
    children: Children,
}

enum Children {
    // 左右の子Nodeを持つ
    Branch { left: Box<Node>, right: Box<Node> },
    // 枝が無い, 親Nodeが終端
    None,
}

Node に統計量 Stats を持たせている。

struct Stats {
    sum: f64,
    sq_sum: f64,
    count: usize,
}

前述したとおり 総和, 二乗総和, 要素数 を持たせる。

Stats の合算

セグメント木においては子ノードが持つ値を合算する処理が重要。というわけで Stats 同士を合算する処理を書く。

impl Stats {
    fn merge(left: &Stats, right: &Stats) -> Stats {
        Stats {
            sum: left.sum + right.sum,
            sq_sum: left.sq_sum + right.sq_sum,
            count: left.count + right.count,
        }
    }
}

見たまんま、それぞれの値を合計して新しい値を作っている。

Stats から平均・分散を計算

先に Stats から平均・分散を計算する処理を書いておこう。

impl Stats {
    fn finalize(&self) -> (f64, f64) {
        let mean = self.sum / self.count as f64;
        let mean_sq = self.sq_sum / self.count as f64;
        let variance = (mean_sq - mean * mean).max(0.0);
        (mean, variance)
    }
}

これも平均や分散の定義を素朴に書き下しただけだ。

セグメント木の構築

配列からセグメント木を構築する処理を書く。
与えられた配列を半分ずつに割っていけばいい。要素が一つなら終端 (children を持たない Node) を作る。複数要素あれば、さらに分割して左右の Node に委譲する。

// 与えられた配列からセグメント木を構築して返す
fn build(data: &[Record]) -> Node {
    Node::build(data, 0, data.len())
}

impl Node {
    fn build(data: &[Record], start: usize, end: usize) -> Node {
        if start + 1 == end {
            // 要素が1つの場合
            let record = &data[start];
            let range_end = data
                .get(end)
                .map_or(record.datetime + chrono::Duration::hours(1), |r| r.datetime);

            Node {
                range: record.datetime..range_end,
                stats: Stats {
                    sum: record.mw,
                    sq_sum: record.mw * record.mw,
                    count: 1,
                },
                children: Children::None,  // 子Nodeは無い
            }
        } else {
            // 要素が複数の場合
            // start, end を真ん中で割って Node::build() を再帰して委譲
            let mid = (start + end) / 2;
            let left = Node::build(data, start, mid);
            let right = Node::build(data, mid, end);

            Node {
                range: left.range.start..right.range.end,
                // 左右の子Nodeのstatsを合算して自身で持っておく
                stats: Stats::merge(&left.stats, &right.stats),
                children: Children::Branch {
                    left: Box::new(left),
                    right: Box::new(right),
                },
            }
        }
    }
}

やっていることはかなりシンプル。
配列インデックス start, end に対して let mid = (start + end) / 2; で真ん中の要素のインデックスを計算し、Node::build() を再帰呼び出しして木の構築を委譲している。
左右の子ノードが出来上がったら、left.stats, right.stats を合算したものを自身で持っておく。

これにて自身のカバーする範囲の集計値を事前計算して各ノードが持っている状態が作れた。

セグメント木への範囲指定問い合わせ

いよいよ範囲を指定して統計量を問い合わせる処理を書こう。

// 指定した範囲の平均・分散を集計する
pub fn query(
    node: &Node,
    query_start: NaiveDateTime,
    query_end_exclusive: NaiveDateTime,
) -> Option<(f64, f64)> {
    node.query(query_start, query_end_exclusive)
        .map(|stats| stats.finalize())
}

impl Node {
    fn query(
        &self,
        query_start: NaiveDateTime,
        query_end_exclusive: NaiveDateTime,
    ) -> Option<Stats> {
        match &self.children {
            Children::Branch { left, right } => {
                // 自身の範囲とクエリ範囲の重複を調べる
                if query_end_exclusive <= self.range.start || self.range.end <= query_start {
                    // 指定された範囲が含まれない場合
                    return None;
                }
                if query_start <= self.range.start && self.range.end <= query_end_exclusive {
                    // 指定された範囲が完全に含まれる場合
                    return Some(self.stats.clone());
                }

                // 指定された範囲を部分的に含む場合
                // 左右の子Nodeに集計を委譲して値を返してもらう
                let left_stats = left.query(query_start, query_end_exclusive);
                let right_stats = right.query(query_start, query_end_exclusive);

                match (left_stats, right_stats) {
                    // 左右両方から値が返っていたら合算する
                    (Some(ls), Some(rs)) => Some(Stats::merge(&ls, &rs)),

                    (Some(ls), None) => Some(ls),
                    (None, Some(rs)) => Some(rs),

                    // 左右どちらからも値が返らないことは起こり得ない
                    (None, None) => unreachable!(),
                }
            }

            Children::None => {
                // 終端の場合は自身がクエリ範囲内に含まれるかだけ確認する
                if query_start <= self.range.start && self.range.start < query_end_exclusive {
                    Some(self.stats.clone())
                } else {
                    None
                }
            }
        }
    }
}

集計処理においても左右の子ノードへの委譲をフル活用する。
実際にやるのは return Some(self.stats.clone()); で自身の持っている集計値を返すか、 return None; で自身の持っている集計値を捨てるかだけ。

実装したセグメント木で集計する

適当な範囲で集計を投げてみる。

fn main() {
    let records = load_records();
    let st = build(&records);

    let start = NaiveDateTime::parse_from_str("2010-01-01 00:00:00", "%Y-%m-%d %H:%M:%S").unwrap();
    let end = NaiveDateTime::parse_from_str("2010-07-01 00:00:00", "%Y-%m-%d %H:%M:%S").unwrap();
    let (mean, variance) = query(&st, start, end).unwrap();

    println!("Mean = {:.2}, Variance = {:.2}", mean, variance);
    // => Mean = 15915.78, Variance = 6883213.05
}

よさそう。

パフォーマンス比較

素朴な線形走査で平均・分散を求める場合とパフォーマンスを比べてみる。

ちなみに「素朴な線形走査」とは、

fn stats(
    records: &[Record],
    query_start: NaiveDateTime,
    query_end_exclusive: NaiveDateTime,
) -> (f64, f64) {
    let mut count = 0;
    let mut sum = 0.0;
    let mut sq_sum = 0.0;

    for record in records {
        if record.datetime < query_start {
            continue;
        }
        if query_end_exclusive <= record.datetime {
            // records がソート済みである前提で早期終了する
            break;
        }
        count += 1;
        sum += record.mw;
        sq_sum += record.mw * record.mw;
    }

    let mean = sum / count as f64;
    let mean_sq = sq_sum / count as f64;
    let variance = (mean_sq - mean * mean).max(0.0);

    (mean, variance)
}

こんな感じ。

単発集計の比較

デモプログラムを実行してみよう。

github.com

$ cargo run demo1
Segment Tree Statistics Demo
Data range: 2004-10-01 01:00:00 ~ 2018-08-03 00:00:00
Enter 'q' to quit

Start (YYYY-MM-DD HH:MM:SS): 2008-01-01 00:00:00
End   (YYYY-MM-DD HH:MM:SS): 2017-01-01 00:00:00

Segment Tree: Mean = 15443.19, Variance = 6774803.89  (20 μs)
Iteration:    Mean = 15443.19, Variance = 6774803.89  (2448 μs)

というわけで8年分, 約70000行を対象に集計をかけた。

実装	平均	分散	実行時間
セグメント木	15443.19	6774803.89	20 μs
線形走査	15443.19	6774803.89	2448 μs

集計結果はもちろん一致。実行時間はセグメント木が120倍早いということになった。
すごい

繰り返し集計する場合の比較

実は単発集計なら素朴に for ループを回して集計したほうがいい。セグメント木は構築自体に時間がかかる。
というわけで、「7日間・28日間のローリングウィンドウで平均・分散が最大になる期間を探す」という問題をやってみる。一つのデータに対して何度も何度も繰り返し問い合わせることになる。これならセグメント木が有利なはずだ。

実際のコードは todays-mitsui/segment-tree-drill/blob/src/main.rs ここにある。

$ cargo run demo2
Rolling Window Statistics Demo
Data range: 2004-10-01 01:00:00 ~ 2018-08-03 00:00:00

=== Window size: 7 days ===
[Segment Tree] (154 ms)
  Max Mean:     20978.69 (at 2007-02-04 11:00:00 ~ 2007-02-11 11:00:00)
  Max Variance: 15282381.82 (at 2010-08-30 12:00:00 ~ 2010-09-06 12:00:00)
[Iteration] (61622 ms)
  Max Mean:     20978.69 (at 2007-02-04 11:00:00 ~ 2007-02-11 11:00:00)
  Max Variance: 15282381.82 (at 2010-08-30 12:00:00 ~ 2010-09-06 12:00:00)

=== Window size: 28 days ===
[Segment Tree] (169 ms)
  Max Mean:     19701.10 (at 2007-01-22 14:00:00 ~ 2007-02-19 14:00:00)
  Max Variance: 11742669.03 (at 2007-08-06 12:00:00 ~ 2007-09-03 12:00:00)
[Iteration] (62046 ms)
  Max Mean:     19701.10 (at 2007-01-22 14:00:00 ~ 2007-02-19 14:00:00)
  Max Variance: 11742669.03 (at 2007-08-06 12:00:00 ~ 2007-09-03 12:00:00)

ウィンドウ幅	実装	実行時間
7日間	セグメント木	154 ms
7日間	線形走査	61622 ms
28日間	セグメント木	169 ms
28日間	線形走査	62046 ms

セグメント木が300倍早いという結果になった。すごすぎる。

まとめ

セグメント木を手で実装して感触を掴んでみた。
シンプルな実装で威力がバツグンなので満足感が高い。

私からは以上です。

補足

もちろんライブラリとして実装している方がいるので、実際に使うなら自前実装よりもライブラリをあたるほうがいいと思う。

• segment-tree - crates.io: Rust Package Registry