ABC340-E Mancala 2

問題へのリンク

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("5 3");
            WillReturn.Add("1 2 3 4 5");
            WillReturn.Add("2 4 0");
            //0 4 2 7 2
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("3 10");
            WillReturn.Add("1000000000 1000000000 1000000000");
            WillReturn.Add("0 1 0 1 0 1 0 1 0 1");
            //104320141 45436840 2850243019
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("1 4");
            WillReturn.Add("1");
            WillReturn.Add("0 0 0 0");
            //1
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static void Main()
    {
        List<string> InputList = GetInputList();
        long[] wkArr = InputList[0].Split(' ').Select(pX => long.Parse(pX)).ToArray();
        long N = wkArr[0];

        long[] AArr = InputList[1].Split(' ').Select(pX => long.Parse(pX)).ToArray();
        long[] BArr = InputList[2].Split(' ').Select(pX => long.Parse(pX)).ToArray();
        long UB = AArr.GetUpperBound(0);

        var IndDualSegmentTree = new DualSegmentTree(UB);
        for (long I = 0; I <= UB; I++) {
            IndDualSegmentTree.RangeAdd(I, I, AArr[I]);
        }

        foreach (long EachB in BArr) {
            long CurrCnt = IndDualSegmentTree.GetVal(EachB);
            IndDualSegmentTree.RangeAdd(EachB, EachB, -CurrCnt);

            // 周回分の加算
            long AddCnt = CurrCnt / N;
            if (AddCnt > 0) {
                IndDualSegmentTree.RangeAdd(0, UB, AddCnt);
            }

            long Rest = CurrCnt % N;
            if (Rest == 0) continue;

            long RangeSta = EachB + 1;
            long RangeEnd = RangeSta + Rest - 1;
            var RangeInfoList = ArrUtil.GetRangeInfoList(RangeSta, RangeEnd, UB);
            foreach (ArrUtil.RangeInfoDef EachRangeInfo in RangeInfoList) {
                IndDualSegmentTree.RangeAdd(EachRangeInfo.IndSta, EachRangeInfo.IndEnd, 1);
            }
        }

        var AnswerList = new List<long>();
        foreach (long EachInd in IndDualSegmentTree.GetNodeIndEnum()) {
            long CurrVal = IndDualSegmentTree.GetVal(EachInd);
            AnswerList.Add(CurrVal);
        }
        Console.WriteLine(LongEnumJoin(" ", AnswerList));
    }

    // セパレータとLong型の列挙を引数として、結合したstringを返す
    static string LongEnumJoin(string pSeparater, IEnumerable<long> pEnum)
    {
        string[] StrArr = Array.ConvertAll(pEnum.ToArray(), pX => pX.ToString());
        return string.Join(pSeparater, StrArr);
    }
}

// 区間加算、1点取得な双対セグ木(フェニック木使用)
#region DualSegmentTree
internal class DualSegmentTree
{
    private long[] mBitArr; // 内部配列(1オリジンなため、添字0は未使用)
    private long mExternalArrUB;

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

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

    // コンストラクタ
    // フェニック木の外部配列(0オリジン)のUBを指定
    internal DualSegmentTree(long pExternalArrUB)
    {
        mExternalArrUB = pExternalArrUB;

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

    // 双対セグメント木の機能
    // 区間加算
    internal void RangeAdd(long pSta, long pEnd, long AddVal)
    {
        pSta++; // 1オリジンに変更
        pEnd++; // 1オリジンに変更

        long ImosSta = pSta;
        long ImosEnd = pEnd + 1;

        // いもす法
        FenwickTree_Add(ImosSta, AddVal);
        if (ImosEnd <= mBitArr.GetUpperBound(0)) {
            FenwickTree_Add(ImosEnd, -AddVal);
        }
    }

    // 双対セグメント木の機能
    // 1点取得
    internal long GetVal(long pInd)
    {
        pInd++; // 1オリジンに変更
        return FenwickTree_GetSum(1, pInd);
    }

    // フェニック木の機能
    // [pSta,pEnd] のSumを返す
    private long FenwickTree_GetSum(long pSta, long pEnd)
    {
        return FenwickTree_GetSum(pEnd) - FenwickTree_GetSum(pSta - 1);
    }

    // フェニック木の機能
    // [0,pEnd] のSumを返す
    private long FenwickTree_GetSum(long pEnd)
    {
        long Sum = 0;
        while (pEnd >= 1) {
            Sum += mBitArr[pEnd];
            pEnd -= pEnd & -pEnd;
        }
        return Sum;
    }

    // フェニック木の機能
    // [I] に Xを加算
    private void FenwickTree_Add(long pI, long pX)
    {
        while (pI <= mBitArr.GetUpperBound(0)) {
            mBitArr[pI] += pX;
            pI += pI & -pI;
        }
    }
}
#endregion

#region ArrUtil
// 配列のユーティリティクラス
internal class ArrUtil
{
    // 区間情報
    internal struct RangeInfoDef
    {
        internal long IndSta;
        internal long IndEnd;
    }

    // IndStaとIndEndとUBを引数とし、
    // IndEndがUBを超えたら、0から開始に変更し、
    // 区間のListを返す
    static internal List<RangeInfoDef> GetRangeInfoList(long pIndSta, long pIndEnd, long UB)
    {
        var WillReturn = new List<RangeInfoDef>();
        if (pIndEnd <= UB) {
            RangeInfoDef WillAdd1;
            WillAdd1.IndSta = pIndSta;
            WillAdd1.IndEnd = pIndEnd;
            WillReturn.Add(WillAdd1);
        }
        else {
            // IndEndがUBを超えてる場合
            RangeInfoDef WillAdd1;
            WillAdd1.IndSta = pIndSta;
            WillAdd1.IndEnd = UB;
            WillReturn.Add(WillAdd1);

            long Rest = pIndEnd - UB;
            RangeInfoDef WillAdd2;
            WillAdd2.IndSta = 0;
            WillAdd2.IndEnd = Rest - 1;
            WillReturn.Add(WillAdd2);
        }
        return WillReturn;
    }
}
#endregion

解説

区間加算の双対セグ木を用意し、
何週できるかと、余りの分の加算をシュミレーションしてます。