何気なく @chokudai さんのタイムラインを見ていたら「セグ木」という言葉が目にとまった。
聞いたことのない言葉だったので調べてみると。
セグ木とは「セグメント木」の略称で、配列の任意の範囲の集計値を効率的に計算できるデータ構造
とのこと。
じゃあ時系列データの統計量の集計に使えるじゃん？と思ったのでやってみる。
ちなみに今回作ったデモプログラムがここにある。
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>,
stats: Stats,
children: Children,
}
enum Children {
Branch { left: Box<Node>, right: Box<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 {
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,
}
} else {
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,
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());
}
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);
}
よさそう。
パフォーマンス比較
素朴な線形走査で平均・分散を求める場合とパフォーマンスを比べてみる。
ちなみに「素朴な線形走査」とは、
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 {
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倍早いという結果になった。すごすぎる。
まとめ
セグメント木を手で実装して感触を掴んでみた。
シンプルな実装で威力がバツグンなので満足感が高い。
私からは以上です。
補足
もちろんライブラリとして実装している方がいるので、実際に使うなら自前実装よりもライブラリをあたるほうがいいと思う。
