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

ABC328-E Modulo MST


問題へのリンク


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 6 328");
            WillReturn.Add("1 2 99");
            WillReturn.Add("1 3 102");
            WillReturn.Add("2 3 86");
            WillReturn.Add("2 4 94");
            WillReturn.Add("2 5 95");
            WillReturn.Add("3 4 81");
            //33
        }
        else if (InputPattern == "Input2") {
            WillReturn.Add("6 5 998244353");
            WillReturn.Add("1 2 337361568");
            WillReturn.Add("1 6 450343304");
            WillReturn.Add("2 3 61477244");
            WillReturn.Add("2 5 745383438");
            WillReturn.Add("4 5 727360840");
            //325437688
        }
        else if (InputPattern == "Input3") {
            WillReturn.Add("8 28 936294041850197");
            WillReturn.Add("1 2 473294720906780");
            WillReturn.Add("1 3 743030800139244");
            WillReturn.Add("1 4 709363019414774");
            WillReturn.Add("1 5 383643612490312");
            WillReturn.Add("1 6 557102781022861");
            WillReturn.Add("1 7 623179288538138");
            WillReturn.Add("1 8 739618599410809");
            WillReturn.Add("2 3 857687812294404");
            WillReturn.Add("2 4 893923168139714");
            WillReturn.Add("2 5 581822471860662");
            WillReturn.Add("2 6 740549363586558");
            WillReturn.Add("2 7 307226438833222");
            WillReturn.Add("2 8 447399029952998");
            WillReturn.Add("3 4 636318083622768");
            WillReturn.Add("3 5 44548707643622");
            WillReturn.Add("3 6 307262781240755");
            WillReturn.Add("3 7 12070267388230");
            WillReturn.Add("3 8 700247263184082");
            WillReturn.Add("4 5 560567890325333");
            WillReturn.Add("4 6 704726113717147");
            WillReturn.Add("4 7 588263818615687");
            WillReturn.Add("4 8 549007536393172");
            WillReturn.Add("5 6 779230871080408");
            WillReturn.Add("5 7 825982583786498");
            WillReturn.Add("5 8 713928998174272");
            WillReturn.Add("6 7 751331074538826");
            WillReturn.Add("6 8 449873635430228");
            WillReturn.Add("7 8 11298381761479");
            //11360716373
        }
        else {
            string wkStr;
            while ((wkStr = Console.ReadLine()) != null) WillReturn.Add(wkStr);
        }
        return WillReturn;
    }

    struct EdgeInfoDef
    {
        internal long FromNode;
        internal long ToNode;
        internal long Cost;
    }
    static List<EdgeInfoDef> mEdgeInfoList = new List<EdgeInfoDef>();

    static void Main()
    {
        List<string> InputList = GetInputList();

        long[] wkArr = { };
        Action<string> SplitAct = pStr =>
            wkArr = pStr.Split(' ').Select(pX => long.Parse(pX)).ToArray();

        SplitAct(InputList[0]);
        long N = wkArr[0];
        long K = wkArr[2];

        foreach (string EachStr in InputList.Skip(1)) {
            SplitAct(EachStr);
            EdgeInfoDef WillAdd;
            WillAdd.FromNode = wkArr[0];
            WillAdd.ToNode = wkArr[1];
            WillAdd.Cost = wkArr[2];
            mEdgeInfoList.Add(WillAdd);
        }

        List<int> IndList = Enumerable.Range(0, mEdgeInfoList.Count).ToList();

        long Answer = long.MaxValue;

        foreach (int[] EachIndArr in RubyPatternClass<int>.Combination(IndList, (int)N - 1)) {
            // 全てのノードが登場するのが必要条件
            var AppearNodeSet = new HashSet<long>();
            foreach (int EachInd in EachIndArr) {
                long FromNode = mEdgeInfoList[EachInd].FromNode;
                long ToNode = mEdgeInfoList[EachInd].ToNode;
                AppearNodeSet.Add(FromNode);
                AppearNodeSet.Add(ToNode);
            }
            if (AppearNodeSet.Count < N) continue;

            // UnionFindで連結かを調べる
            var InsUnionFind = new UnionFind();
            for (long I = 1; I <= N; I++) {
                InsUnionFind.MakeSet(I);
            }

            foreach (int EachInd in EachIndArr) {
                long FromNode = mEdgeInfoList[EachInd].FromNode;
                long ToNode = mEdgeInfoList[EachInd].ToNode;
                InsUnionFind.Unite(FromNode, ToNode);
            }

            // 代表ノードが1つだけなら連結している
            var RootSet = new HashSet<long>();
            for (long I = 1; I <= N; I++) {
                RootSet.Add(InsUnionFind.FindSet(I));
            }
            if (RootSet.Count == 1) {
                long CurrCost = 0;

                foreach (int EachInd in EachIndArr) {
                    CurrCost += mEdgeInfoList[EachInd].Cost;
                    CurrCost %= K;
                }
                Answer = Math.Min(Answer, CurrCost);
            }
        }
        Console.WriteLine(Answer);
    }
}

#region RubyPatternClass
// Rubyの場合の数
internal static class RubyPatternClass<Type>
{
    // 組合せを返す
    private struct JyoutaiDef_Combination
    {
        internal int CurrP;
        internal List<int> SelectedIndList;
    }
    internal static IEnumerable<Type[]> Combination(IEnumerable<Type> pEnum, int pR)
    {
        if (pR == 0) yield break;
        Type[] pArr = pEnum.ToArray();
        if (pArr.Length < pR) yield break;

        var Stk = new Stack<JyoutaiDef_Combination>();
        JyoutaiDef_Combination WillPush;
        for (int I = pArr.GetUpperBound(0); 0 <= I; I--) {
            WillPush.CurrP = I;
            WillPush.SelectedIndList = new List<int>() { I };
            Stk.Push(WillPush);
        }

        while (Stk.Count > 0) {
            JyoutaiDef_Combination Popped = Stk.Pop();

            // クリア判定
            if (Popped.SelectedIndList.Count == pR) {
                var WillReturn = new List<Type>();
                Popped.SelectedIndList.ForEach(X => WillReturn.Add(pArr[X]));
                yield return WillReturn.ToArray();
                continue;
            }

            for (int I = pArr.GetUpperBound(0); Popped.CurrP + 1 <= I; I--) {
                WillPush.CurrP = I;
                WillPush.SelectedIndList = new List<int>(Popped.SelectedIndList) { I };
                Stk.Push(WillPush);
            }
        }
    }
}
#endregion

#region UnionFind
// UnionFindクラス
internal class UnionFind
{
    private class NodeInfoDef
    {
        internal long ParentNode;
        internal long Rank;
    }
    private Dictionary<long, NodeInfoDef> mNodeInfoDict =
        new Dictionary<long, NodeInfoDef>();

    // 要素が1つである木を森に追加
    internal void MakeSet(long pNode)
    {
        NodeInfoDef WillAdd = new NodeInfoDef();
        WillAdd.ParentNode = pNode;
        WillAdd.Rank = 0;
        mNodeInfoDict[pNode] = WillAdd;
    }

    // 合併処理
    internal void Unite(long pX, long pY)
    {
        long XNode = FindSet(pX);
        long YNode = FindSet(pY);
        long XRank = mNodeInfoDict[XNode].Rank;
        long YRank = mNodeInfoDict[YNode].Rank;

        if (XRank > YRank) {
            mNodeInfoDict[YNode].ParentNode = XNode;
        }
        else {
            mNodeInfoDict[XNode].ParentNode = YNode;
            if (XRank == YRank) {
                mNodeInfoDict[YNode].Rank++;
            }
        }
    }

    // ノードを引数として、木の根を取得
    internal long FindSet(long pTargetNode)
    {
        // 根までの経路上のノードのList
        var PathNodeList = new List<long>();

        long CurrNode = pTargetNode;
        while (CurrNode != mNodeInfoDict[CurrNode].ParentNode) {
            PathNodeList.Add(CurrNode);
            CurrNode = mNodeInfoDict[CurrNode].ParentNode;
        }

        // 経路圧縮 (親ポインタの付け替え)
        foreach (long EachPathNode in PathNodeList) {
            mNodeInfoDict[EachPathNode].ParentNode = CurrNode;
        }
        return CurrNode;
    }

    internal void DebugPrint()
    {
        foreach (var EachPair in mNodeInfoDict.OrderBy(pX => pX.Key)) {
            Console.WriteLine("mNodeInfoDict[{0}].ParentNode={1}",
                EachPair.Key, EachPair.Value.ParentNode);
        }
    }
}
#endregion


解説

全域木は木なので
枝数は、(ノード数-1)になります。

なので
ノード数 C 枝数
なだけの枝の組み合わせを列挙し、
UnionFindで連結なら解候補にしてます。