Hatena::Grouptopcoder

kuuso1の日記

2014-11-13

CodeForces #277 div2 D. Valid Sets

| 01:56 | CodeForces #277 div2 D. Valid Sets - kuuso1の日記 を含むブックマーク はてなブックマーク - CodeForces #277 div2 D. Valid Sets - kuuso1の日記

http://codeforces.com/contest/486/problem/D

  • n個のノードを持つtreeと、整数d が与えられる
  • それぞれのノード i は値 a[ i ] を持つ (i= 1 to n)
  • 以下のような集合 S をValid と呼ぶ
    • Sは空でない
    • Sは連結である
    • Max{A[k]| k in S} - Min{A[k]| k in S} ≤d
  • このようなValid Setの総数を mod 1000000007 で求める

(制約)1 ≤ n ≤ 2000, 1 ≤ d ≤ 2000,1 ≤ A[i] ≤ 2000 (i= 1 to n)

方針:各ノードを根として、根が最大値となるような極大部分木を作ることを考える。

そのような部分木が求まれば、その根を共有する部分木は全てValid Setとなる。

従って、根から出発して、根以下であって差がd以内のノードを順次繋いでいけばよい。

ここで、根と同じ値に出会った場合は、重複をなくすため、そのノードが過去に根として計上されているか否かで、繋ぐかどうかを判断すればよい。(Point2)

また、ある部分木に対して、根を共有する部分木の個数の数え上げは、子ノードの選び方の組み合わせになるため、

各子ノードの値が計算できていれば、Point1の要領で計算することができる。

f:id:kuuso1:20141114014905p:image

コードとしては、極大木を作る部分と数え上げの部分は並行してdfsで行うことができるため、各ノードについて1回のdfsで求めることができる。

計算量はO(n^2)。コードは以下

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

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

class Sol{
	public void Solve(){
		mod=(long)(1e9+7);
		
		long Ans=0;
		for(int i=0;i<N;i++){
			Ans+=dfs(i,-1,i);
			Ans%=mod;
		}
		
		Console.WriteLine(Ans);
		
	}
	
	long dfs(int now,int parent,int root){
		long ret=1;
		foreach(int n in E[now]){
			if(n!=parent){
				if(A[n]<A[root] && A[root]-A[n]<=D){
					ret*=(1+dfs(n,now,root));
					ret%=mod;
					continue;
				}
				if(A[n]==A[root] && n>root){
					ret*=(1+dfs(n,now,root));
					ret%=mod;
					continue;
				}
			}
		}
		return ret;
	}
	
	long mod;
	
	int N;
	int D;
	int[] A;
	List<int>[] E;

	public Sol(){
		var d=ria();
		D=d[0];N=d[1];
		A=ria();
		E=new List<int>[N];
		for(int i=0;i<N;i++){
			E[i]=new List<int>();
		}
		for(int i=0;i<N-1;i++){
			var dd=ria();
			E[dd[0]-1].Add(dd[1]-1);//0-origin
			E[dd[1]-1].Add(dd[0]-1);//0-origin
		}
	}

	static int[] ria(){return Array.ConvertAll(Console.ReadLine().Split(' '),e=>int.Parse(e));}
}

コンテスト中は各ノードを最大値とみなして繋ぐという発想がなかったため解けなかったが、twitterでの諸氏のヒントをもとに一晩考えると簡潔なコードで書けたので満足。