Hackerrank Recurrent on a Tree

最新推荐文章于 2018-03-01 17:43:49 发布

原创最新推荐文章于 2018-03-01 17:43:49 发布 · 373 阅读

1 ·

CC 4.0 BY-SA版权

Hackerrank 专栏收录该内容

4 篇文章

订阅专栏

本文介绍了一种结合树形动态规划与矩阵快速幂技术的算法，用于解决一类特定的树形结构上的路径和问题。该算法通过巧妙地利用矩阵操作实现了高效的递归过程，避免了传统递归方法中的重复计算。

题目

定义一条路径上节点值为x ，它的价值就为fib[x]，求一棵树的所有路径的和

搞一搞矩阵操作就可以DP了

#include<stdio.h>
#include<string.h>
#include<iostream>
#include<algorithm>
using namespace std;
typedef long long ll;
const ll MOD=(ll)1e9+7;
const int MAXN=(int)2e5+10;
struct Matrix{
	ll m[3][3];
	void init(){memset(m,0,sizeof(m));}
}one,fib;
Matrix mul(const Matrix &a,const Matrix &b){
	Matrix re;
	re.init();
	for(int i=1;i<=2;i++)for(int j=1;j<=2;j++){
		for(int k=1;k<=2;k++){
			re.m[i][j]=(re.m[i][j]+a.m[i][k]*b.m[k][j])%MOD;
		}
	}
	return re;
}
Matrix add(const Matrix &a,const Matrix &b){
	Matrix re;
	for(int i=1;i<=2;i++)for(int j=1;j<=2;j++){
		re.m[i][j]=(a.m[i][j]+b.m[i][j])%MOD;
	}
	return re;
}
Matrix power(Matrix a,ll b){
	Matrix re=one;
	while(b){
		if(b&1)re=mul(re,a);
		a=mul(a,a);
		b>>=1;
	}
	return re;
}

struct edge{
	int to,nxt;
}ed[MAXN<<1];
int head[MAXN],cnt;
ll c[MAXN];
void addedge(int u,int v){
	ed[cnt].to=v;
	ed[cnt].nxt=head[u];
	head[u]=cnt++;
}
Matrix a1,a2,dp[MAXN],vl[MAXN];
void dfs(int u,int pre){
//	vl[u]=dp[u]=power(fib,c[u]);    //这里会爆栈，写到main里面才A的
	a1=add(a1,vl[u]);
	for(int i=head[u];i!=-1;i=ed[i].nxt){
		int v=ed[i].to;
		if(v!=pre){
			dfs(v,u);
			a2=add(a2,mul(dp[u],dp[v]));
			dp[u]=add(dp[u],mul(dp[v],vl[u]));
		}
	}
}
int main()
{
	memset(head,-1,sizeof(head));
	a1.init();
	a2.init();
	one.init();
	fib.m[1][1]=0;
	fib.m[1][2]=fib.m[2][1]=fib.m[2][2]=1;	
	one.m[1][1]=one.m[2][2]=1;
	int n;
	scanf("%d",&n);
	for(int i=1;i<n;i++){
		int u,v;
		scanf("%d%d",&u,&v);
		addedge(u,v);
		addedge(v,u);
	}
	for(int i=1;i<=n;i++)scanf("%lld",&c[i]),vl[i]=power(fib,c[i]),dp[i]=vl[i];
	dfs(1,0);
	printf("%lld\n",(a1.m[1][1]+a1.m[1][2]+a2.m[1][1]*2+a2.m[1][2]*2)%MOD);
	return 0;
}