kuuso1の日記

Week of Code - 19 Coolguy and Two Subsequences

HackerRank, 数学, UnionFind | 03:56 |

https://www.hackerrank.com/contests/w19/challenges/coolguy-and-two-subsequences

• N 個の整数 A[n] 1 ≤ n ≤ N が与えられる
• 次の疑似コードで与えられる数値を mod 1e9+7 で求めよ

（制約）1≤ N ≤200000,　1≤ A[n] ≤ 10^9

//f(a, b) is a function that returns the minimum element in interval [a, b]

ans = 0

for a -> [1, n]
    for b -> [a, n]
        for c -> [b + 1, n]
            for d -> [c, n]
                ans = ans + min(f(a, b), f(c, d))

・疑似コードの意味するところは、全ての閉区間のペアに関してその最小値の和を求めよという問題。

・ナイーブに書くとO(N^5)くらいになるのでとてもとても。

・例えばA[n]全体の最小値を含む区間が選ばれている時は、相手が何であっても最小値がansに足されるので、いくつそのような組み合わせがあるかを考えるのがよさそう。

・A[n]を大きい順に見ていくことで、自分が最小値になる組み合わせがいくつあるかを数えていく。

・図中のz(z+1)/2やw(w+1)/2の部分については、（一度登場するとあとはずっと登場するので）毎回計算する必要はなく和を覚えておけばよい。

　ただし、zやwの値や連結成分の個数が下図のように変わっていくので、UnionFindで管理しながら適宜足し引きしながら更新する。

　コードでは、着目している数値に連結してしまう成分の Length(Length+1)/2 を総和から引いておき、ループの最後に更新して加えている

・A[n]の値が重複している場合については、どちらかでカウントすれば問題ないため、ソートした順に真に大小がついていると考えてよい。

コードは以下

using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Text;

class TEST{
	static void Main(){
		Sol mySol =new Sol();
		mySol.Solve();
	}
}

class Sol{
	public void Solve(){
		
		Pair[] P=new Pair[N];
		for(int i=0;i<N;i++){
			P[i]=new Pair(i,A[i]);
		}
		Array.Sort(P,(x,y)=>-x.Val.CompareTo(y.Val));
		int[] Ord=new int[N];
		int[] Idx=new int[N];
		for(int i=0;i<N;i++){
			Ord[P[i].Idx]=i;
			Idx[i]=P[i].Idx;
		}
		
		var UF=new UnionFind(N);
		long mod=(long)1e9+7;
		long ans=0;
		long sqsum=0;
		
		for(int ord=0;ord<N;ord++){
			int idx=Idx[ord];
			long pat=0;
			
			if(idx>0 && Ord[idx-1]<ord){
				long sq0=(long)(UF.Right(idx-1)-UF.Left(idx-1)+1);
				sq0=sq0*(sq0+1)/2;
				sq0%=mod;
				sqsum-=sq0;sqsum%=mod;if(sqsum<0)sqsum+=mod;
				UF.Union(idx,idx-1);
			}
			if(idx<N-1 && Ord[idx+1]<ord){
				long sq0=(long)(UF.Right(idx+1)-UF.Left(idx+1)+1);
				sq0=sq0*(sq0+1)/2;
				sq0%=mod;
				sqsum-=sq0;sqsum%=mod;if(sqsum<0)sqsum+=mod;
				UF.Union(idx+1,idx);
			}
			
			// **WWW**XXAYYY***ZZZ***
			
			long x=(long)(idx-UF.Left(idx));
			long y=(long)(UF.Right(idx)-idx);
			long home1=(x+1)*(y+1);
			if(home1>=mod)home1%=mod;
			
			// case(1)    XXAYYY >I * I <ZZZ,WWW
			pat+=(sqsum*home1)%mod;
			if(pat>=mod)pat%=mod;
			
			// case(2)    XXAYYY > I1,I2 　　　 
			long patx=x*(x+1)*(x+2)/6;
			patx%=mod;
			patx*=(y+1);patx%=mod;
			long paty=y*(y+1)*(y+2)/6;
			paty%=mod;
			paty*=(x+1);paty%=mod;
			
			pat+=(patx+paty)%mod;
			pat%=mod;
			
			ans+=(A[idx]*pat)%mod;
			ans%=mod;
			
			long sq1=(long)(UF.Right(idx)-UF.Left(idx)+1);
			sq1=sq1*(sq1+1)/2;
			sq1%=mod;
			sqsum+=sq1;
			sqsum%=mod;
		}
		
		Console.WriteLine(ans);
		
	}
	
	class Pair{
		public int Idx;
		public long Val;
		public Pair(int idx,long val){
			Val=val;
			Idx=idx;
		}
	}
	
	int N;
	long[] A;
	public Sol(){
		N=ri();
		//A=rla();
		var ss=rs().Split(new char[]{' '},StringSplitOptions.RemoveEmptyEntries);
		A=new long[N];
		for(int i=0;i<N;i++)A[i]=long.Parse(ss[i]);
	}

	static String rs(){return Console.ReadLine();}
	static int ri(){return int.Parse(Console.ReadLine());}
	static long rl(){return long.Parse(Console.ReadLine());}
	static double rd(){return double.Parse(Console.ReadLine());}
	static String[] rsa(char sep=' '){return Console.ReadLine().Split(sep);}
	static int[] ria(char sep=' '){return Array.ConvertAll(Console.ReadLine().Split(sep),e=>int.Parse(e));}
	static long[] rla(char sep=' '){return Array.ConvertAll(Console.ReadLine().Split(sep),e=>long.Parse(e));}
	static double[] rda(char sep=' '){return Array.ConvertAll(Console.ReadLine().Split(sep),e=>double.Parse(e));}
}

class UnionFind{
	
	int[] parent;
	int[] mem;
	int compo;
	int N;
	int[] L,R; //Leftest,Rightest of the component
	
	public UnionFind(int n_){
		Initialize(n_);
	}
	
	public void Initialize(int n_){
		N=n_;
		parent=new int[N];
		mem=new int[N];
		L=new int[N];
		R=new int[N];
		for(int i=0;i<N;i++){
			parent[i]=i;
			mem[i]=1;
			L[i]=i;
			R[i]=i;
		}
		compo=N;
	}
	
	public int Parent(int a){
		if(parent[a]==a)return a;
		return parent[a]=Parent(parent[a]);
	}
	
	public bool areSameGp(int a,int b){
		return Parent(a)==Parent(b);
	}
	
	public void Union(int a,int b){
		a=Parent(a);b=Parent(b);
		if(a==b)return;
		parent[a]=b;
		mem[b]+=mem[a];
		compo-=1;
		int ll=Math.Min(L[a],L[b]);
		int rr=Math.Max(R[a],R[b]);
		L[b]=L[a]=ll;
		R[b]=R[a]=rr ;
	}
	
	public int Left(int x){
		return L[Parent(x)];
	}

	public int Right(int x){
		return R[Parent(x)];
	}
}

・最初、各数値の左右に自分より大きいのがどれだけいるかを全列挙することも考えたんですが（セグ木でO(logN*logN*N）くらいでやる）、

　散在する区間群を管理するのにUnionFindもいるなぁと思ってよくよく考えると

@sigma425さんのCompetitive Programming Advent Calendar 2015の記事でも紹介されていたUnionFindに値を持たせるテクが使える！っていうのでかなりテンション上がったよね。

ツイートする