SegmentTree.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
 
namespace MyCaffe.trainers.common
{
    public class SegmentTree
    {
        protected int m_nCapacity;
        protected OPERATION m_op;
        protected double[] m_rgfValues;
 
        public enum OPERATION
        {
            SUM,
            MIN
        }
 
        public SegmentTree(int nCapacity, OPERATION oper, float fNeutralElement)
        {
            if (nCapacity <= 0 || (nCapacity % 2) != 0)
                throw new Exception("The capacity must be positive and a power of 2.");
 
            m_nCapacity = nCapacity;
            m_op = oper;
            m_rgfValues = new double[2 * nCapacity];
 
            for (int i = 0; i < m_rgfValues.Length; i++)
            {
                m_rgfValues[i] = fNeutralElement;
            }
        }
 
        private double operation(double f1, double f2)
        {
            switch (m_op)
            {
                case OPERATION.MIN:
                    return Math.Min(f1, f2);
 
                case OPERATION.SUM:
                    return f1 + f2;
 
                default:
                    throw new Exception("Unknown operation '" + m_op.ToString() + "'!");
            }
        }
 
        private double reduce_helper(int nStart, int nEnd, int nNode, int nNodeStart, int nNodeEnd)
        {
            if (nEnd < 0)
            {
                Trace.WriteLine("**SegmentTree: nEnd is less than zero and should not be.");
                return m_rgfValues[nNode];
            }
 
            if (nStart == nNodeStart && nEnd == nNodeEnd)
                return m_rgfValues[nNode];
 
            int nMid = (nNodeStart + nNodeEnd) / 2;
 
            if (nEnd <= nMid)
            {
                return reduce_helper(nStart, nEnd, 2 * nNode, nNodeStart, nMid);
            }
            else
            {
                if (nMid + 1 <= nStart)
                {
                    return reduce_helper(nStart, nMid, 2 * nNode + 1, nMid + 1, nNodeEnd);
                }
                else
                {
                    double f1 = reduce_helper(nStart, nMid, 2 * nNode, nNodeStart, nMid);
                    double f2 = reduce_helper(nMid + 1, nEnd, 2 * nNode + 1, nMid + 1, nNodeEnd);
                    return operation(f1, f2);
                }
            }
        }
 
        public double reduce(int nStart, int? nEnd = null)
        {
            if (!nEnd.HasValue)
                nEnd = m_nCapacity;
 
            if (nEnd < 0)
                nEnd += m_nCapacity;
 
            nEnd -= 1;
 
            return reduce_helper(nStart, nEnd.Value, 1, 0, m_nCapacity - 1);
        }
 
        public double this[int nIdx]
        {
            get
            {
                if (nIdx < 0 || nIdx >= m_nCapacity)
                    throw new Exception("The index is out of range!");
 
                return m_rgfValues[m_nCapacity + nIdx];
            }
            set
            {
                nIdx += m_nCapacity;
                m_rgfValues[nIdx] = value;
 
                nIdx /= 2;
 
                while (nIdx >= 1)
                {
                    m_rgfValues[nIdx] = operation(m_rgfValues[2 * nIdx], m_rgfValues[2 * nIdx + 1]);
                    nIdx /= 2;
                }
            }
        }
    }
 
    public class SumSegmentTree : SegmentTree
    {
        public SumSegmentTree(int nCapacity)
            : base(nCapacity, OPERATION.SUM, 0.0f)
        {
        }
 
        public double sum(int nStart = 0, int? nEnd = null)
        {
            return reduce(nStart, nEnd);
        }
 
        public int find_prefixsum_idx(double fPrefixSum)
        {
            if (fPrefixSum < 0)
                throw new Exception("The prefix sum must be greater than zero.");
 
            double fSum = sum() + 1e-5;
            if (fPrefixSum > fSum)
                throw new Exception("The prefix sum cannot exceed the overall sum of '" + fSum.ToString() + "'!");
 
            int nIdx = 1;
            while (nIdx < m_nCapacity) // while non-leaf
            {
                if (m_rgfValues[2 * nIdx] > fPrefixSum)
                {
                    nIdx = 2 * nIdx;
                }
                else
                {
                    fPrefixSum -= m_rgfValues[2 * nIdx];
                    nIdx = 2 * nIdx + 1;
                }
            }
 
            return nIdx - m_nCapacity;
        }
    }
 
    public class MinSegmentTree : SegmentTree
    {
        public MinSegmentTree(int nCapacity)
            : base(nCapacity, OPERATION.MIN, float.MaxValue)
        {
        }
 
        public double min(int nStart = 0, int? nEnd = null)
        {
            return reduce(nStart, nEnd);
        }
    }
}