AtCoderのABC    次のABCの問題へ    前のABCの問題へ

ABC369-C Count Arithmetic Subarrays


問題へのリンク


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("4");
            WillReturn.Add("3 6 9 3");
            //8
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("5");
            WillReturn.Add("1 1 1 1 1");
            //15
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("8");
            WillReturn.Add("87 42 64 86 72 58 44 30");
            //22
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

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

        if (AArr.Length == 1) {
            Console.WriteLine(1);
            return;
        }

        // 階差数列を作る
        var KaisaList = new List<long>();
        for (long I = 1; I <= UB; I++) {
            KaisaList.Add(AArr[I] - AArr[I - 1]);
        }
        long[] KaisaArr = KaisaList.ToArray();

        List<RunLenLongArr.RunLenLongArrInfo> RunLenLongArrInfoList =
            RunLenLongArr.DeriveRunLenLongArrInfoList(KaisaArr);

        long Answer = AArr.Length;
        foreach (RunLenLongArr.RunLenLongArrInfo RunLenLongInfo in RunLenLongArrInfoList) {
            long Cnt = RunLenLongInfo.RunLen;

            long AddVal = Cnt * (Cnt + 1) / 2;
            Answer += AddVal;
        }
        Console.WriteLine(Answer);
    }
}

#region RunLen
// ランレングス圧縮(longの配列)
internal class RunLenLongArr
{
    // ランレングス圧縮情報
    internal struct RunLenLongArrInfo
    {
        internal long AppearNum;
        internal long RunLen;
    }

    // ランレングス圧縮結果を返す
    internal static List<RunLenLongArrInfo> DeriveRunLenLongArrInfoList(long[] pArr)
    {
        var WillReturn = new List<RunLenLongArrInfo>();

        // 空配列の対応
        if (pArr.Length == 0) {
            return WillReturn;
        }

        long PrevNum = pArr[0];
        long SeqLen = 0;
        for (int I = 0; I <= pArr.GetUpperBound(0); I++) {
            if (pArr[I] != PrevNum) {
                RunLenLongArrInfo WillAdd;
                WillAdd.AppearNum = PrevNum;
                WillAdd.RunLen = SeqLen;
                WillReturn.Add(WillAdd);
                SeqLen = 0;
                PrevNum = pArr[I];
            }
            SeqLen++;

            if (I == pArr.GetUpperBound(0)) {
                RunLenLongArrInfo WillAdd;
                WillAdd.AppearNum = pArr[I];
                WillAdd.RunLen = SeqLen;
                WillReturn.Add(WillAdd);
            }
        }
        return WillReturn;
    }
}
#endregion


解説

1 2 3 6 9 12 24 36 48 60 90
のような数列で階差数列を考えると
 1 1 1 3 3 12 12 12 12 30
になります。
大小関係は意味がないのでアルファベットにします。
A A A B B C C C C D
すると
連続したAからは、(3 + 2 + 1) 通り
連続したBからは、(2 + 1) 通り
連続したCからは、(4 + 3 + 2 + 1) 通り
連続したDからは、1 通り

後は、1からnまでの自然数の和の公式を使えば
項数が2以上の等差数列の個数を高速に数え上げできます。

そして、
項数が1の等差数列も解に計上すれば、解を求めることができます。