題目
找一個點當作root，把樹吊起來。
我們可以發現，基本上對於非root的點，只要不跟連接parent的邊同顏色即可。
所以就在dfs多帶一個參數，去避免同色狀況即可！
AC Code
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
|
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> ii;
vector<vector<ii>> g;
vector<int> ans;
int mx = 0;
void dfs(int u, int p, int usedColor)
{
int idx = 1;
for (auto cur : g[u]) {
int v = cur.first;
int edgeNumber = cur.second;
if (v == p) // skip parent
continue;
if (usedColor == idx)
idx++;
mx = max(mx, idx);
ans[edgeNumber] = idx;
dfs(v, u, idx);
idx++;
}
}
int main()
{
int n;
scanf("%d", &n);
g = vector<vector<ii>>(n, vector<ii>());
ans = vector<int>(n - 1, -1);
for (int i = 0; i < n - 1; i++) {
int u, v;
scanf("%d %d", &u, &v);
u--;
v--;
g[u].push_back(ii(v, i));
g[v].push_back(ii(u, i));
}
dfs(0, -1, -1);
printf("%d\n", mx);
for (int i = 0; i < n - 1; i++)
printf("%d\n", ans[i]);
return 0;
}
|
真的是被數學打敗。
這一題只要先把所求用數學寫出來，之後移項一下，就可以變成簡單枚舉題目。
令 $S_i = A_0 + A_1 + A_2 …$ 且 $S_0 = 0$，就標準的 prefix sum 的形式
我們的所求寫為數學算式的話 $(S_j - S_i) \pmod k = (j - i)$
移項一下得到 $(S_j - j) \pmod k = (S_i - i) \pmod k$ 且 $(j - i) < k$
之後就很單純啦！建立一個 map 來做 counting。我們只要把所有的 $(S_j - j) \pmod k$ 都依照順序枚舉出來，記得踢掉不符合 $(j - i) < k$ 的 $(S_j - j) \pmod k$ 就是啦！
小心 long long…
AC Code
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
|
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> ii;
int main()
{
int n, k;
scanf("%d %d", &n, &k);
ll pre[n + 1] = {0};
for (int i = 1; i <= n; i++) {
ll num;
scanf("%lld", &num);
pre[i] = pre[i - 1] + num;
}
map<ll, int> cnt;
ll ans = 0;
for (int j = 0; j <= n; j++) { // start with 0...
if (j - k >= 0)
cnt[(pre[j - k] - (j - k) + k) % k]--;
ans += cnt[(pre[j] - j + k) % k];
cnt[(pre[j] - j + k) % k]++;
}
printf("%lld\n", ans);
return 0;
}
|