ABC452-E You WILL Like 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("6 4");
            WillReturn.Add("1 6 9 2 3 1");
            WillReturn.Add("1 10 3 7");
            //508
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("20 20");
            WillReturn.Add("36625 195265 98908 111868 111868 47382 147644 472464 472464 416653 111868 195265 327972 327972 262769 75439 381156 451275 36625 195265");
            WillReturn.Add("327972 111868 416653 177330 340019 262769 47382 262769 47382 340019 47382 262769 327972 327972 359676 381156 327972 36625 451275 381156");
            //58141644
        }
        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();
    }

    const long Hou = 998244353;

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

        long[] AArr = GetSplitArr(InputList[1]);
        long[] BArr = GetSplitArr(InputList[2]);

        var ADict = new Dictionary<long, long>();
        for (long I = 0; I <= AArr.GetUpperBound(0); I++) {
            ADict[I + 1] = AArr[I];
        }

        var BDict = new Dictionary<long, long>();
        for (long I = 0; I <= BArr.GetUpperBound(0); I++) {
            BDict[I + 1] = BArr[I];
        }

        var ProdDict = new Dictionary<long, long>();
        long ProdSum = 0;
        for (long I = 1; I <= N; I++) {
            ProdDict[I] = ADict[I];
            ProdDict[I] %= Hou;
            ProdSum += ADict[I] * I;
            ProdSum %= Hou;
        }

        var InsFenwick_Tree = new Fenwick_Tree(N, Hou);
        for (long I = 1; I <= N; I++) {
            InsFenwick_Tree[I] = ADict[I];
        }

        long Answer = 0;

        // Jを固定して考える
        for (long J = 1; J <= M; J++) {
            long CalcResult = ProdSum;

            long RangeSta = J;
            long RangeEnd = RangeSta + J - 1;
            long Omomi = 1;
            while (true) {
                if (RangeSta > N) break;
                RangeEnd = Math.Min(RangeEnd, N);

                long RangeSum = InsFenwick_Tree.GetSum(RangeSta, RangeEnd);
                long wkProd = RangeSum * Omomi;
                wkProd %= Hou;
                wkProd *= J;
                wkProd %= Hou;
                CalcResult -= wkProd;
                CalcResult %= Hou;
                if (CalcResult < 0) {
                    CalcResult += Hou;
                }
                Omomi++;

                RangeSta += J;
                RangeEnd += J;
            }
            long CurrAnswer = BDict[J] * CalcResult;
            CurrAnswer %= Hou;
            Answer += CurrAnswer;
            Answer %= Hou;
        }
        Console.WriteLine(Answer);
    }
}

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

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

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

    // コンストラクタ
    internal Fenwick_Tree(long pExternalArrUB, long pHou)
    {
        mExternalArrUB = pExternalArrUB;

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

        mHou = pHou;
    }

    // インデクサ
    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)
    {
        long Result = GetSum(pEnd) - GetSum(pSta - 1);

        Result %= mHou;
        if (Result < 0) Result += mHou;
        return Result;
    }

    // [0,pEnd] のSumを返す
    internal long GetSum(long pEnd)
    {
        pEnd++; // 1オリジンに変更

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

    // [I] に Xを加算
    internal void Add(long pI, long pX)
    {
        pI++; // 1オリジンに変更

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

解説

Σ(I=1からN)Σ(J=1からM)(AI * BJ * (I % J))
でJを固定して考えると
BJ * Σ(I=1からN)(AI * (I % J))
となります。
具体例で、J=5 , N=100で考えます。
BJ * Σ(I=1から100)(AI * (I % 5))
Σ(I=1から100)(AI * (I % 5)) <= Σ(I=1から100)(AI * I)
ですので
Σ(I=1から100)(AI * I) を仮の値として、余計な分を引き算することを考えます。

I= 5だと A5  *  5 - A5  * 5 * 1 (商が1のグループ)
I= 6だと A6  *  6 - A6  * 5 * 1 (商が1のグループ)
I= 7だと A7  *  7 - A7  * 5 * 1 (商が1のグループ)
I= 8だと A8  *  8 - A8  * 5 * 1 (商が1のグループ)
I= 9だと A9  *  9 - A9  * 5 * 1 (商が1のグループ)
I=10だと A10 * 10 - A10 * 5 * 2 (商が2のグループ)
I=11だと A11 * 11 - A11 * 5 * 2 (商が2のグループ)
I=12だと A12 * 12 - A12 * 5 * 2 (商が2のグループ)
I=13だと A13 * 13 - A13 * 5 * 2 (商が2のグループ)
I=14だと A14 * 14 - A14 * 5 * 2 (商が2のグループ)
I=15だと A15 * 15 - A15 * 5 * 3 (商が3のグループ)
I=16だと A16 * 16 - A16 * 5 * 3 (商が3のグループ)
I=17だと A17 * 17 - A17 * 5 * 3 (商が3のグループ)
I=18だと A18 * 18 - A18 * 5 * 3 (商が3のグループ)
となるので
AI * I の 総和を求めておき
AI のフェニック木を使い、商ごとのAIの合計から求めた値を求め、
引き算すれば良いと分かります。