Skip to content

Instantly share code, notes, and snippets.

@ibako
Last active February 22, 2022 13:33
Show Gist options
  • Save ibako/61726d85306447a46913c089eb38c23e to your computer and use it in GitHub Desktop.
Save ibako/61726d85306447a46913c089eb38c23e to your computer and use it in GitHub Desktop.
Segment Tree with lazy evaluation
using System;
public class SegmentTreeLazy<T> : SegmentTreeLazy<T, T>
{
public SegmentTreeLazy(T[] data, T identity,
Func<T, T, T> dataChildrenToParent, Func<T, T, T> lazyToData, Func<T, T, T> lazyToDataLeaf,
Func<T, T, T> lazyToLazy, Func<T, int, T> lazySegmentSizeToLazy)
: base(data, identity, identity, dataChildrenToParent, lazyToData, lazyToDataLeaf, lazyToLazy, lazySegmentSizeToLazy)
{
}
}
public class SegmentTreeLazy<TData, TLazy>
{
private readonly int _N;
private readonly TData[] _Data;
private readonly TLazy[] _Lazy;
private readonly TData _DataIdentity;
private readonly TLazy _LazyIdentity;
private readonly Func<TData, TData, TData> _DataChildrenToParent;
private readonly Func<TData, TLazy, TData> _LazyToData;
private readonly Func<TData, TLazy, TData> _LazyToDataLeaf;
private readonly Func<TLazy, TLazy, TLazy> _LazyToLazy;
private readonly Func<TLazy, int, TLazy> _LazySegmentSizeToLazy;
public SegmentTreeLazy(TData[] data, TData dataIdentity, TLazy lazyIdentity,
Func<TData, TData, TData> dataChildrenToParent, Func<TData, TLazy, TData> lazyToData, Func<TData, TLazy, TData> lazyToDataLeaf,
Func<TLazy, TLazy, TLazy> lazyToLazy, Func<TLazy, int, TLazy> lazySegmentSizeToLazy)
{
// data.Lengthを2進数表記したときの長さ
var len = Convert.ToString(data.Length - 1, 2).Length;
this._N = (int)Math.Pow(2, len);
this._Data = new TData[2 * this._N - 1];
this._Lazy = new TLazy[2 * this._N - 1];
this._DataIdentity = dataIdentity;
this._LazyIdentity = lazyIdentity;
this._DataChildrenToParent = dataChildrenToParent;
this._LazyToData = lazyToData;
this._LazyToDataLeaf = lazyToDataLeaf;
this._LazyToLazy = lazyToLazy;
this._LazySegmentSizeToLazy = lazySegmentSizeToLazy;
for (var i = 0; i < this._Lazy.Length; i++) this._Lazy[i] = lazyIdentity;
for (var i = this._N - 1; i < 2 * this._N - 1; i++)
{
var idx = i - (this._N - 1);
this._Data[i] = idx >= data.Length
? dataIdentity
: data[idx];
}
for (var i = this._N - 2; i >= 0; i--)
{
this._Data[i] = dataChildrenToParent(this._Data[2 * i + 1], this._Data[2 * i + 2]);
}
}
public void Update(int index, TLazy value)
{
this.Update(index, index + 1, value);
}
// 区間 [startIndex, endIndex) をまとめて更新する(遅延評価)
public void Update(int startIndex, int endIndex, TLazy value)
{
this.UpdateImpl(startIndex, endIndex, value, 0, 0, this._N);
}
private void UpdateImpl(int startIndex, int endIndex, TLazy value, int i, int l, int r)
{
this.Eval(i, r - l);
if (startIndex <= l && r <= endIndex)
{
// 現在のセグメントを全部使う
this._Lazy[i] = this._LazyToLazy(this._Lazy[i], value);
this.Eval(i, r - l);
}
else if (startIndex < r && l < endIndex)
{
// 現在のセグメントの一部を使う
this.UpdateImpl(startIndex, endIndex, value, i * 2 + 1, l, (l + r) / 2);
this.UpdateImpl(startIndex, endIndex, value, i * 2 + 2, (l + r) / 2, r);
this._Data[i] = this._DataChildrenToParent(this._Data[i * 2 + 1], this._Data[i * 2 + 2]);
}
}
public TData Query(int index)
{
return this.Query(index, index + 1);
}
public TData Query(int startIndex, int endIndex)
{
return this.QueryImpl(startIndex, endIndex, 0, 0, this._N);
}
private TData QueryImpl(int startIndex, int endIndex, int i, int l, int r)
{
this.Eval(i, r - l);
if (r <= startIndex || endIndex <= l)
{
// 現在のセグメントが問い合わせの範囲外
return this._DataIdentity;
}
else if(startIndex <= l && r <= endIndex)
{
// 現在のセグメントを全部使う
return this._Data[i];
}
else
{
// 現在のセグメントの一部を使う
var vl = this.QueryImpl(startIndex, endIndex, i * 2 + 1, l, (l + r) / 2);
var vr = this.QueryImpl(startIndex, endIndex, i * 2 + 2, (l + r) / 2, r);
return this._DataChildrenToParent(vl, vr);
}
}
private void Eval(int i, int size)
{
if (this._Lazy[i].Equals(this._LazyIdentity)) return;
if (i < this._N - 1)
{
// 葉でなければ子に遅延評価値を伝播
this._Lazy[i * 2 + 1] = this._LazyToLazy(this._Lazy[i * 2 + 1], this._Lazy[i]);
this._Lazy[i * 2 + 2] = this._LazyToLazy(this._Lazy[i * 2 + 2], this._Lazy[i]);
this._Data[i] = this._LazyToData(this._Data[i], this._LazySegmentSizeToLazy(this._Lazy[i], size));
}
else
{
this._Data[i] = this._LazyToDataLeaf(this._Data[i], this._LazySegmentSizeToLazy(this._Lazy[i], size));
}
this._Lazy[i] = this._LazyIdentity;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment