ABC353-E Yet Another 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("3");
            WillReturn.Add("ab abc arc");
            //4
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("11");
            WillReturn.Add("ab bb aaa bba baba babb aaaba aabbb a a b");
            //32
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    static void Main()
    {
        List<string> InputList = GetInputList();
        string[] SArr = InputList[1].Split(' ');

        var CntDict = new Dictionary<decimal, long>();
        foreach (string EachS in SArr) {
            var InsRollingHash = new RollingHash(EachS);
            for (int I = 0; I <= EachS.Length - 1; I++) {
                decimal CurrHash = InsRollingHash.GetRangeHash(0, I);
                if (CntDict.ContainsKey(CurrHash) == false) {
                    CntDict[CurrHash] = 0;
                }
                CntDict[CurrHash]++;
            }
        }

        long Answer = 0;
        foreach (long EachCnt in CntDict.Values) {
            Answer += EachCnt * (EachCnt - 1) / 2;
        }
        Console.WriteLine(Answer);
    }
}

#region RollingHash
// ローリングハッシュ
internal class RollingHash
{
    const decimal Base = 100M;
    const decimal Hou = 67280421310721M;

    // ハッシュ値[終了Ind]で、[0,終了Ind]のハッシュ値を保持
    internal decimal[] mHashArr;

    // 桁の重みの配列
    decimal[] mOmomiArr;

    // コンストラクタ
    internal RollingHash(string pTargetStr)
    {
        decimal Omomi = 1;
        long UB = pTargetStr.Length - 1;

        mHashArr = new decimal[UB + 1];
        mOmomiArr = new decimal[UB + 1];
        for (int I = 0; I <= UB; I++) {
            mOmomiArr[I] = Omomi;
            if (I > 0) {
                mHashArr[I] += mHashArr[I - 1] * Base;
                mHashArr[I] %= Hou;
            }
            mHashArr[I] += pTargetStr[I];
            mHashArr[I] %= Hou;
            Omomi *= Base;
            Omomi %= Hou;
        }
    }

    // [Sta,End]のハッシュ値を返す
    internal decimal GetRangeHash(long pSta, long pEnd)
    {
        decimal WillReturn = mHashArr[pEnd];
        if (pSta > 0) {
            long Range = pEnd - pSta + 1;
            decimal PrevVal = mHashArr[pSta - 1];
            decimal MinusVal = PrevVal * mOmomiArr[Range];
            MinusVal %= Hou;
            WillReturn -= MinusVal;
            if (WillReturn < 0) {
                WillReturn += Hou;
            }
        }
        return WillReturn;
    }
}
#endregion

解説

a
aa
ab
ab
abc
bcd
bcd
c
c
というサンプルで考えると、

a   5件
aa  1件
ab  3件
abc 1件
b   2件
bc  2件
bcd 2件
c 2件

長さ3の接頭辞が一致するペア
例えば、abcは
aで1回
abで1回
abcで1回
として、合計3とカウントしても良いです。

なので、接頭辞ごとの
件数 choose 2
の総和で解けるので、ローリングハッシュを使えば良いです。