ABC241-E Putting Candies

問題へのリンク

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("2 1 6 3 1");
            //11
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("10 1000000000000");
            WillReturn.Add("260522 914575 436426 979445 648772 690081 933447 190629 703497 47202");
            //826617499998784056
        }
        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 K = wkArr[1];

        long[] AArr = InputList[1].Split(' ').Select(pX => long.Parse(pX)).ToArray();

        var AppearList = new List<long>();
        var AppearSet = new HashSet<long>();

        long[] PreCycleArr;
        long[] CycleArr;

        long CurrInd = 0;
        while (true) {
            if (AppearSet.Contains(CurrInd)) {
                PreCycleArr = AppearList.TakeWhile(pX => pX != CurrInd).ToArray();
                CycleArr = AppearList.SkipWhile(pX => pX != CurrInd).ToArray();
                break;
            }
            AppearList.Add(CurrInd);
            AppearSet.Add(CurrInd);

            CurrInd += AArr[CurrInd];
            CurrInd %= AArr.Length;
        }

        // 周期に入る前の分を加算
        long RestN = K;
        long TakeCnt = Math.Min(RestN, PreCycleArr.Length);
        long Answer = PreCycleArr.Take((int)TakeCnt).Sum(pX => AArr[pX]);
        RestN -= PreCycleArr.Length;
        if (RestN == 0) {
            Console.WriteLine(Answer);
            return;
        }

        // 周期の分を加算
        if (CycleArr.Length <= RestN) {
            long OneCycleSum = CycleArr.Sum(pX => AArr[pX]);
            Answer += OneCycleSum * (RestN / CycleArr.Length);
            RestN %= CycleArr.Length;
        }

        // 残りの分を加算
        Answer += CycleArr.Take((int)RestN).Sum(pX => AArr[pX]);

        Console.WriteLine(Answer);
    }
}

解説

数列が、途中からサイクルになることに着目してます。

類題

ABC179-E Sequence Sum