连通块计数

链接：https://www.nowcoder.com/acm/contest/204/I
来源：牛客网

题目描述

小 A 有一棵长的很奇怪的树，他由 n 条链和 1 个点作为根构成，第 i 条链有 ai 个点，每一条链的一端都与根结点相连。
现在小 A 想知道，这棵长得奇怪的树有多少非空的连通子树，你只需要输出答案对 998244353 取模的值即可

输入描述:

第一行一个正整数 n
第二行 n 个正整数 a1…an

输出描述:

输出答案对 998244353 取模后的值

示例1

输入

2
1 1

输出

6

备注:

1≤ n≤ 105
1≤ ai≤ 107

先贴代码

#include <iostream>
#define ll long long
const ll mod = 998244353;
using namespace std;
int main()
{
	int n;
	scanf("%d", &n);
	ll s1 = 0, s2 = 1, s3 = 0;
	for (int i = 0; i < n; i++)
	{
		scanf("%lld", &s1);
		s2 += s1 * s2;
		s2 %= mod;
		s3 += s1 * (s1 + 1) / 2;
		s3 %= mod;
	}
	printf("%lld\n", (s2 + s3) % mod);
}

自己的思路

挺有意思的一个题目，记录一下，s2表示当前链到达根节点或其他链节点的连通子树的个数，s3表示当前链的内部可以构成的连通子树的个数，不理解不妨到纸上模拟一下，这个是真的妙啊。