はまやんはまやんはまやん

hamayanhamayan's blog

Mishka and trip [Codeforces 365 : B]

問題

http://codeforces.com/contest/703/problem/B

n個の都市があり、そのうちk個が首都である
首都間は以下のルールで辺が作られる

  • 隣り合う都市間には辺がある。都市iと都市i+1間には辺がある(都市nと都市1間にもある)
  • 首都は全ての頂点間と辺がある

都市にコストciが掛かっており、辺はその両端の都市のコストの積である
このとき、全ての辺のコストの和を求めよ

3 <= n <= 10^5
1 <= k <= n
1 <= ci <= 10^4

考察

1. 愚直解を考えてみる
2. 全ての頂点間の辺について考えてコストを足せばいい
3. これだとO(n^2)

4. もうちょっと賢い解法を考えてみる
5. 今回は「首都は全ての頂点間と辺がある」条件において計算量を減らせそう
6. とある首都iと他の頂点との間の辺のコストの総和を考えてみる

以下、ci = c[i]と書きます

c[1]*c[i] + c[2]*c[i] + ... + c[i-1]*c[i] + c[i+1]*c[i] + ... c[n] * c[i]
= c[i] (c[1] + c[2] + ... + c[i-1] + c[i+1] + ... + c[n])
= c[i] (sumall - c[i])
※ sumall = c[1] + c[2] + ... + c[n]
<||

7. という訳である首都と他の頂点間の辺のコストの総和は c[i](sumall - c[i])
8. これは使えそうだな(という仮定を得ました)

9. 全首都に対してこれを計算して総和を取ると、どちらも首都である辺をダブって数えている
10. なので、全首都間のコストの総和を考えてみます。

>||
m = 4として考えてみる。
全首都が id1, id2, id3, id4 とする
c[id1]*c[id2] + c[id1]*c[id3] + c[id1]*c[id4] + c[id2]*c[id3] + c[id2]*c[id4] + c[id3]*c[id4]
= c[id1](c[id2] + c[id3] + c[id4]) + c[id2](c[id3] + c[id4]) + c[id3]*c[id4]
こう考えると、最初に全てのコストを足しておき、c[id1]を引いて残りと掛ける、次にc[id2]を引いて残りと掛けるを繰り返せば行けそう

11. これを使えば首都間のコストの総和が得られるので、これを引きます
12. あとは環状のループのうち、どちらも首都じゃないものを足して解です

13. これを書いてて思ったのが、全部の辺のコストの総和からどちらも首都じゃない辺のコストを引けば単に答えになるのでは・・・?

実装

http://codeforces.com/contest/703/submission/19628155

typedef long long ll;
int n, k;
int c[101010];
bool cap[101010];
//-----------------------------------------------------------------
int main() {
	cin >> n >> k;
	rep(i, 0, n) scanf("%d", &c[i]);
	rep(i, 0, k) {
		int id; scanf("%d", &id);
		cap[id - 1] = true;
	}

	ll ans = 0;

	// 考察5~8
	ll sum = 0;
	rep(i, 0, n) sum += c[i];
	rep(i, 0, n) if (cap[i]) ans += (sum - c[i]) * c[i];

	// 考察12
	rep(i, 0, n) {
		int j = (i + 1) % n;
		if (!cap[i] && !cap[j]) ans += c[i] * c[j];
	}

	// 考察9~11
	sum = 0;
	rep(i, 0, n) if (cap[i]) sum += c[i];
	rep(i, 0, n) if (cap[i]) {
		sum -= c[i];
		ans -= sum * c[i];
	}

	cout << ans << endl;
}