ABC146-E Rem of Sum is Num

問題へのリンク

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 4");
            WillReturn.Add("1 4 2 3 5");
            //4
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("8 4");
            WillReturn.Add("4 2 4 2 4 2 4 2");
            //7
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("10 7");
            WillReturn.Add("14 15 92 65 35 89 79 32 38 46");
            //8
        }
        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();

        long Result = Solve(K, AArr);
        Console.WriteLine(Result);
    }

    static long Solve(long pK, long[] pAArr)
    {
        long UB = pAArr.GetUpperBound(0);

        // K=1の場合
        if (pK == 1) {
            return pAArr.Count(pX => pX % pK == 1);
        }

        // 1引いてから、Kを法とした余りに変更する
        for (long I = 0; I <= UB; I++) {
            pAArr[I]--;
            if (pAArr[I] < 0) {
                pAArr[I] += pK;
            }
            pAArr[I] %= pK;
        }

        // 累積和を求める
        long[] RunSumArr = new long[UB + 1];
        for (int I = 0; I <= UB; I++) {
            RunSumArr[I] += pAArr[I];
            if (I > 0) RunSumArr[I] += RunSumArr[I - 1];
            RunSumArr[I] %= pK;
        }

        // 件数[mod]なDict
        var DosuuDict = new Dictionary<long, long>();

        long Answer = 0;

        // 区間の右端を全探索
        for (long I = 0; I <= UB; I++) {
            long RemoveInd = I - pK;
            if (RemoveInd >= 0) {
                long RemoveVal = RunSumArr[RemoveInd];
                DosuuDict[RemoveVal]--;
            }

            long CurrVal = RunSumArr[I];
            if (DosuuDict.ContainsKey(CurrVal) == false) {
                DosuuDict[CurrVal] = 0;
            }
            DosuuDict[CurrVal]++;

            long CurrPettern = DosuuDict[CurrVal] - 1;
            if (I < pK - 1 && CurrVal == 0) CurrPettern++;

            // Console.WriteLine("右端Ind={0}での件数={1}", I, CurrPettern);

            Answer += CurrPettern;
        }
        return Answer;
    }
}

解説

対象区間になった数列が
満たすべき条件を考えると
Sum(配列の値) % K = 配列の個数
で、これは
配列の個数がK以上なら解なしで
配列の個数がK未満なら、
SQLっぽく書くと
Sum(Val) % K = Count(*)で これは
Sum(Val) % K = Sum(1)で 位項すると
Sum(Val - 1) % K = 0
となります。

後は、(Val - 1) の累積和を求めておいて、
区間の和が0になる、区間の数が解になりますので、

累積和の度数分布表を、Dictで管理しながら
区間の右端を全探索してます。