ABC378-E Mod Sigma Problem

問題へのリンク

C#のソース

using System;
using System.Collections.Generic;
using System.Linq;

class Program
{
    static string InputPattern = "InputX";

    static List<string> GetInputList()
    {
        var WillReturn = new List<string>();

        if (InputPattern == "Input1") {
            WillReturn.Add("3 4");
            WillReturn.Add("2 5 0");
            //10
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("10 100");
            WillReturn.Add("320 578 244 604 145 839 156 857 556 400");
            //2736
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static long[] GetSplitArr(string pStr)
    {
        return (pStr == "" ? new string[0] : pStr.Split(' ')).Select(pX => long.Parse(pX)).ToArray();
    }

    static void Main()
    {
        List<string> InputList = GetInputList();
        long[] wkArr = GetSplitArr(InputList[0]);
        long M = wkArr[1];

        long[] AArr = GetSplitArr(InputList[1]);
        long UB = AArr.GetUpperBound(0);

        // 累積和を求める
        long[] RunSumArr = (long[])AArr.Clone();
        for (long I = 0; I <= UB; I++) {
            if (I > 0) {
                RunSumArr[I] += RunSumArr[I - 1];
            }
            RunSumArr[I] %= M;
        }

        // 個数[mod]なフェニック木
        var Ins_Fenwick_Tree1 = new Fenwick_Tree(M);

        // 総和[mod]なフェニック木
        var Ins_Fenwick_Tree2 = new Fenwick_Tree(M);

        foreach (long EachVal in RunSumArr) {
            Ins_Fenwick_Tree1[EachVal] += 1;
            Ins_Fenwick_Tree2[EachVal] += EachVal;
        }

        long Answer = 0;

        // 区間の左端を全探索
        for (long I = 0; I <= UB; I++) {
            long PrevRunSum = 0;
            if (0 < I) {
                PrevRunSum = RunSumArr[I - 1];
            }

            long RangeSta1 = PrevRunSum;
            long RangeEnd1 = M;

            // PrevSum以上の集計
            long RangeSum1 = Ins_Fenwick_Tree2.GetSum(RangeSta1, RangeEnd1);
            RangeSum1 -= PrevRunSum * Ins_Fenwick_Tree1.GetSum(RangeSta1, RangeEnd1);

            // PrevSum未満の集計
            long RangeSum2 = 0;
            if (PrevRunSum > 0) {
                long RangeSta2 = 0;
                long RangeEnd2 = PrevRunSum - 1;

                long PairVal = M - PrevRunSum;

                RangeSum2 = Ins_Fenwick_Tree2.GetSum(RangeSta2, RangeEnd2);
                RangeSum2 += PairVal * Ins_Fenwick_Tree1.GetSum(RangeSta2, RangeEnd2);
            }

            Answer += RangeSum1 + RangeSum2;

            // フェニック木を差分更新
            long CurrRunSum = RunSumArr[I];
            Ins_Fenwick_Tree1[CurrRunSum] -= 1;
            Ins_Fenwick_Tree2[CurrRunSum] -= CurrRunSum;
        }
        Console.WriteLine(Answer);
    }
}

// フェニック木
#region Fenwick_Tree
internal class Fenwick_Tree
{
    private long[] mBitArr;
    private long mExternalArrUB;

    // ノードのIndexの列挙を返す
    internal IEnumerable<long> GetNodeIndEnum()
    {
        for (long I = 0; I <= mExternalArrUB; I++) {
            yield return I;
        }
    }

    // 木のノードのUBを返す
    internal long GetUB()
    {
        return mExternalArrUB;
    }

    // コンストラクタ(外部配列のUBのみ指定)
    internal Fenwick_Tree(long pExternalArrUB)
    {
        mExternalArrUB = pExternalArrUB;

        // フェニック木の外部配列は0オリジンで、
        // フェニック木の内部配列は1オリジンなため、2を足す
        mBitArr = new long[pExternalArrUB + 2];
    }

    // コンストラクタ(初期化用の配列指定)
    internal Fenwick_Tree(long[] pArr)
        : this(pArr.GetUpperBound(0))
    {
        for (long I = 0; I <= pArr.GetUpperBound(0); I++) {
            this.Add(I, pArr[I]);
        }
    }

    // コンストラクタ(初期化用のList指定)
    internal Fenwick_Tree(List<long> pList)
        : this(pList.Count - 1)
    {
        for (int I = 0; I <= pList.Count - 1; I++) {
            this.Add(I, pList[I]);
        }
    }

    // Indのチェック
    private void IndCheck(long pInd)
    {
        if (pInd < 0) throw new Exception("pInd < 0");
        if (mExternalArrUB < pInd) throw new Exception("UB < pInd");
    }

    // Indの大小チェック
    private void IndRangeCheck(long pSta, long pEnd)
    {
        IndCheck(pSta);
        IndCheck(pEnd);
        if (pSta > pEnd) throw new Exception("pSta > pEnd");
    }

    // インデクサ
    internal long this[long pInd]
    {
        get { return GetSum(pInd, pInd); }
        set { Add(pInd, value - GetSum(pInd, pInd)); }
    }

    // [pSta,pEnd] のSumを返す
    internal long GetSum(long pSta, long pEnd)
    {
        IndRangeCheck(pSta, pEnd);

        long Result = GetSum(pEnd);
        if (pSta > 0) {
            Result -= GetSum(pSta - 1);
        }
        return Result;
    }

    // [0,pEnd] のSumを返す
    internal long GetSum(long pEnd)
    {
        IndCheck(pEnd);

        pEnd++; // 1オリジンに変更

        long Sum = 0;
        while (pEnd >= 1) {
            Sum += mBitArr[pEnd];
            pEnd -= pEnd & -pEnd;
        }
        return Sum;
    }

    // [I] に Xを加算
    internal void Add(long pI, long pX)
    {
        IndCheck(pI);

        pI++; // 1オリジンに変更

        while (pI <= mBitArr.GetUpperBound(0)) {
            mBitArr[pI] += pX;
            pI += pI & -pI;
        }
    }
}
#endregion

解説

11 10
3 1 4 1 5 9 2 6 5 3 5
で考えます。

まずmod10での累積和を求めます
3 1 4 1 5 9 2 6 5 3 5
3 4 8 9 4 3 5 1 6 9 4

そして、区間の左端を全探索することを考えます。
初回は、累積和をフェニック木に入れておけば、総和を簡単に求めることができます。

次に左端を1つ右にずらした時を考えます。
左端の3が消えるので、3以上の数は、3引けばよく
3未満の数は、法が10なので、+10-3で+7となります。
3 4 8 9 4 3 5 1 6 9 4
↓
  1 5 6 1 0 2 8 3 5 1

この状態での総和は、
最初の累積和で
個数[mod]なフェニック木と
総和[mod]なフェニック木を作っておけば、
3以上と3未満で分けて考え、高速に求めることができます。

左端が進んだ時も、フェニック木を差分更新し、
modでの総和を、高速に求めることができます。