POJ 3370 - Halloween treats

http://poj.org/problem?id=3370

概要

\(n\) 軒の家があり，\(i\) 番目の家では \(a_i\) 個の飴をもらえる． どの家に行けば \(c\) 人で等分できる数の飴を手に入れられるか答える． ただし，答えが複数存在するときはそのうちのどれを答えてもよく (Special Judge)， 答えが存在しないときは「no sweets」と出力する．

制約

• \(1 \le c \le n \le 10 ^ 5\)
• \(1 \le a_i \le 10 ^ 5\)

解法

飴の数に関してはすべて \(\mbox{mod } c\) で考えればよい．

部分和 \(s_i = a_1 + a_2 + \cdots + a_i\) を考える． \(s_i = 0\) ならば \(1, 2, \cdots, i\) が答えである． そうでないとき，\(s_i\) のとりうる値は高々 \(c-1\) 通りしか存在しない． ここで制約から \(c-1 < n\) なので，鳩の巣原理より \(i < j \land s_i = s_j\) なる \(i, j\) が存在する． このとき \(a_{i+1} + a_{i+2} + \cdots + a_j = 0\) となるので，これを答えればよい． したがって「no sweets」となることはない．

これで \(\mathrm{O}(n)\) で解けるわけだけど，出力が多くて上記のアルゴリズムを愚直に実装すると TLE する… 複数答えが存在するときはなるべく短くなるやつを選ぶようにしてみたら 1000MS ちょっとで通った．

 1 #include <cstdio>
 2 #include <algorithm>
 3 using namespace std;
 4 
 5 int main()
 6 {
 7   int C, N;
 8   while (scanf("%d %d", &C, &N) != EOF && C != 0) {
 9     static const int K = 100000;
10     static int prev[K];
11     fill_n(prev, N, -1);
12     int b = -1, e = -1;
13     int acc = 0;
14     for (int i = 0; i < N; i++) {
15       int x;
16       scanf("%d", &x);
17       acc = (acc + x) % C;
18       if (x % C == 0) {
19         b = e = i;
20       } else if (acc == 0) {
21         if (b == -1 || i < e-b) {
22           b = 0;
23           e = i;
24         }
25       } else if (prev[acc] != -1) {
26         if (b == -1 || i-prev[acc]-1 < e-b) {
27           b = prev[acc]+1;
28           e = i;
29         }
30       } else {
31         prev[acc] = i;
32       }
33     }
34     printf("%d", b+1);
35     for (int i = b+1; i <= e; i++) {
36       printf(" %d", i+1);
37     }
38     putchar('\n');
39   }
40   return 0;
41 }