Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <bits/stdc++.h>
- typedef int64_t Int;
- const Int mod = (Int)1e9+7;
- // Sum of arithmetic progression: 1 + 2 + ... + last = last * (last+1) / 2
- Int sum(Int last) {
- return last * (last+1) % mod * ((mod+1) / 2) % mod;
- }
- // Calculate sum of walks between each pair of vertices in caterpillar tree:
- Int solve(const std::vector<Int> & d) {
- Int answer = 0, n = (Int)d.size();
- // Calculate sum of walks between each pair of main vertices:
- for (Int last = n-1; last > 0; --last) {
- answer = (answer + sum(last)) % mod;
- }
- // Calculate sum of walks between each pair from main and additional vertices:
- for (Int i = 0; i < n; ++i) {
- answer = (answer + sum(i+1) * d[i] % mod) % mod;
- answer = (answer + sum(n-i) * d[i] % mod) % mod;
- }
- answer = (answer + mod - std::accumulate(d.begin(), d.end(), Int(0))) % mod; // substract repetitive walks
- // Calculate sum of walks between each pair of additional vertices in one group:
- for (Int i = 0; i < n; ++i) {
- answer = (answer + d[i] * (d[i]-1) % mod) % mod;
- }
- // Calculate suffix sums over vector d:
- // suffix1[i] = d[i] + d[i+1] + d[i+2] + ... + d[n-1]
- // suffix2[i] = (i+1) * d[i] + (i+2) * d[i+1] + (i+3) * d[i+2] + ... + n * d[n-1]
- std::vector<Int> suffix1(n+1, 0), suffix2(n+1, 0);
- for (Int i = n-1; i >= 0; --i) {
- suffix1[i] = (suffix1[i+1] + d[i]) % mod;
- suffix2[i] = (suffix2[i+1] + d[i] * (i+1)) % mod;
- }
- // Calculate walks between each pair of vertices from different groups:
- for (Int i = 0; i < n; ++i) {
- Int mul = (suffix2[i+1] + (mod + 1 - i) * suffix1[i+1] % mod) % mod;
- answer = (answer + d[i] * mul % mod) % mod;
- }
- return answer;
- }
- int main() {
- std::ios_base::sync_with_stdio(false); std::cin.tie(0); std::cout.tie(0);
- Int n; std::cin >> n;
- std::vector<Int> d(n); for (auto& it : d) std::cin >> it;
- std::cout << solve(d) << std::endl;
- return 0;
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement