TCO 2010 Finals Hard

With no better idea of spending some time on the Las Vegas - New York plane, I've decided to try coding the hard problem from the TCO 2010 finals (statement linked). This took me 26 minutes and I got 618 points in the practice room — knowing the solution in advance and using IDEA instead of a simple text editor (rng_58's time in the finals was 48 minutes, for 463 points).

I hope the code with the below explanations can serve as a semi-editorial while we're waiting for the official editorial.

  1: import java.util.*; 
  2:  
  3: public class ColorfulJewelry { 
  4:     static final long MOD = 1000000009; 
  5:     int height; 
  6:     int width; 
  7:     int maxCount; 
  8:     String[] jewels; 
  9:  
 10:     class RectangleSums { 
 11:         int[][] sums; 
 12:  
 13:         public RectangleSums(char ch) { 
 14:             sums = new int[height + 1][width + 1]; 
 15:             for (int r = 0; r < height; ++r) { 
 16:                 int colSum = 0; 
 17:                 for (int c = 0; c < width; ++c) { 
 18:                     if (jewels[r].charAt(c) == ch) 
 19:                         ++colSum; 
 20:                     sums[r + 1][c + 1] = colSum + sums[r][c + 1]; 
 21:                 } 
 22:             } 
 23:         } 
 24:  
 25:         int getSum(int r1, int c1, int r2, int c2) { 
 26:             if (r1 >= r2 || c1 >= c2) 
 27:                 return 0; 
 28:             return get(r2, c2) + get(r1, c1) - get(r1, c2) - get(r2, c1); 
 29:         } 
 30:  
 31:         private int get(int r, int c) { 
 32:             if (r < 0) 
 33:                 r = 0; 
 34:             if (c < 0) 
 35:                 c = 0; 
 36:             if (r > height) 
 37:                 r = height; 
 38:             if (c > width) 
 39:                 c = width; 
 40:             return sums[r][c]; 
 41:         } 
 42:     } 
 43:  
 44:     public int getChains(String[] jewels, int k) { 
 45:         this.jewels = jewels; 
 46:         height = jewels.length; 
 47:         width = jewels[0].length(); 
 48:         RectangleSums[] sums = new RectangleSums[3]; 
 49:         for (int i = 0; i < 3; ++i) 
 50:             sums[i] = new RectangleSums("RGB".charAt(i)); 
 51:         maxCount = height * width; 
 52:         int[][] maxCounts = new int[maxCount + 1][maxCount + 1]; 
 53:         for (int i = 0; i <= maxCount; ++i) 
 54:             Arrays.fill(maxCounts[i], -1); 
 55:         for (int r1 = 0; r1 < 1 || r1 + k <= height; ++r1) { 
 56:             for (int c1 = 0; c1 < 1 || c1 + k <= width; ++c1) { 
 57:                 for (int r2 = r1; r2 < r1 + 1 || r2 + k <= height; ++r2) { 
 58:                     for (int c2 = 0; c2 < 1 || c2 + k <= width; ++c2) { 
 59:                         int[] count = new int[3]; 
 60:                         for (int i = 0; i < 3; ++i) { 
 61:                             count[i] += sums[i].getSum(r1, c1, r1 + k, c1 + k); 
 62:                             count[i] += sums[i].getSum(r2, c2, r2 + k, c2 + k); 
 63:                             count[i] -= sums[i].getSum(Math.max(r1, r2), Math.max(c1, c2), Math.min(r1, r2) + k, Math.min(c1, c2) + k); 
 64:                         } 
 65:                         maxCounts[count[0]][count[1]] = Math.max(maxCounts[count[0]][count[1]], count[2]); 
 66:                     } 
 67:                 } 
 68:             } 
 69:         } 
 70:         for (int r = maxCount; r >= 0; --r) 
 71:             for (int c = maxCount; c >= 0; --c) { 
 72:                 if (r > 0) 
 73:                     maxCounts[r - 1][c] = Math.max(maxCounts[r - 1][c], maxCounts[r][c]); 
 74:                 if (c > 0) 
 75:                     maxCounts[r][c - 1] = Math.max(maxCounts[r][c - 1], maxCounts[r][c]); 
 76:             } 
 77:         prepare(); 
 78:         long res = 0; 
 79:         for (int r = 0; r <= maxCount; ++r) 
 80:             for (int g = 0; g <= maxCount; ++g) 
 81:                 if (maxCounts[r][g] >= 0) { 
 82:                     res = (res + solve(r, g, maxCounts[r][g])) % MOD; 
 83:                 } 
 84:         // Divide by 2
 85:         res = (res * ((MOD + 1) / 2)) % MOD; 
 86:         // Subtract 1
 87:         res = (res + MOD - 1) % MOD; 
 88:         return (int) res; 
 89:     } 
 90:  
 91:     int[][] combinations; 
 92:     int[][] combinationSums; 
 93:  
 94:     private void prepare() { 
 95:         combinations = new int[maxCount + 1][maxCount + 1]; 
 96:         combinations[0][0] = 1; 
 97:         for (int n = 1; n <= maxCount; ++n) { 
 98:             combinations[n][0] = 1; 
 99:             for (int i = 1; i <= maxCount; ++i) 
100:                 combinations[n][i] = (int) (((long) combinations[n - 1][i] + combinations[n - 1][i - 1]) % MOD); 
101:         } 
102:         combinationSums = new int[maxCount + 1][maxCount + 1]; 
103:         for (int rplusg = 0; rplusg <= maxCount; ++rplusg) { 
104:             combinationSums[rplusg][0] = combinations[rplusg][0]; 
105:             for (int maxb = 1; maxb + rplusg <= maxCount; ++maxb) { 
106:                 combinationSums[rplusg][maxb] = (int) (((long) combinationSums[rplusg][maxb - 1] + combinations[rplusg + maxb][maxb]) % MOD); 
107:             } 
108:         } 
109:     } 
110:  
111:     private long solveSimple(int r, int g, int maxB) { 
112:         long simple = combinations[r + g][r]; 
113:         simple *= combinationSums[r + g][maxB]; 
114:         simple %= MOD; 
115:         return simple; 
116:     } 
117:  
118:  
119:     private long solve(int r, int g, int maxB) { 
120:         long simple = solveSimple(r, g, maxB); 
121:         // Now to the symmetric ones
122:         long symmetric = 0; 
123:         if (r % 2 != 0) { 
124:             if (g % 2 != 0) { 
125:                 // This can't be symmetric
126:             } else { 
127:                 symmetric = solveSimple(r / 2, g / 2, maxB / 2); 
128:             } 
129:         } else if (g % 2 != 0) { 
130:             symmetric = solveSimple(r / 2, g / 2, maxB / 2); 
131:         } else { 
132:             // Even number of Bs
133:             symmetric = solveSimple(r / 2, g / 2, maxB / 2); 
134:             if (maxB > 0) { 
135:                 // Odd number of Bs
136:                 symmetric += solveSimple(r / 2, g / 2, (maxB - 1) / 2); 
137:                 symmetric %= MOD; 
138:             } 
139:         } 
140:         return (simple + symmetric) % MOD; 
141:     } 
142: } 

The entry point method, getChains, is located in lines 44-89. First, we create 3 RectangleSums objects each tracking one kind of jewel. The object itself, located in lines 10-42, keeps track of sums for each (0,0)-(x,y) rectangle, and that allows it to calculate the sum for an arbitrary rectangle using inclusion-exclusion.

Lines 51-69 are responsible for trying all possible combinations of two squares. Please note the loop boundaries: although the problem statement allowed for the squares to be outside the field, it is not really necessary since we can always use only some of the jewels we get. The only case when this is necessary is when k is larger than the size of the field - that's the reason for the "r1 < 1 ||" clause and similar ones for other loops. We also start r2 from r1, not from 0, to cut the total number of cases to consider in half.

ACRush has told me that his hard has timed out because he lacked the above optimizations, that's why I already had them in my mind. During the actual contest, I guess only extensive testing could lead to being so careful here. The solution runs for 1.4s on the worst case, and there's only one such case in the test data (but I guess any test case with a 44x44 field and k=1 will present this running time).

Lines 61-63 calculate the total counts of each jewel in both rectangles using inclusion-exclusion. Line 65 stores the result in the following manner: for every pair of #R and #G we've seen, we store the maximum #B. Since we can always discard some jewels, the maximum is enough for our purposes.

Lines 70-76 "spread" the maximum towards lower values of #R and #G. By that we calculate the following: what is the maximum number of Bs that can appear on a chain with exactly the given amount of Rs and Gs?

Method prepare() calculates two arrays that I've mentioned in the live coverage and that are later used to quickly count chains: combinations[n][k] is the number of ways to choose k items out of n, and combinationSums[rplusg][maxb] is equal to combinations[rplusg][0] + combinations[rplusg+1][1]+...+combinations[rplusg+maxb][maxb].

Lines 78-83 call method solve() which answers the following question: "how many chains are there with exactly R red jewels, exactly G green jewels, and between 0 and maxB blue jewels?".

"how many" is defined in a strange way, though. In order to calculate the number of chains up to reversal, we need to calculate the total number of chains, then the number of symmetric chains, and then divide all non-symmetric ones by 2, yielding the formula (total-symmetric)/2+symmetric=(total+symmetric)/2. We make solve() return the value of "total+symmetric", and divide by 2 in line 85 of the main method. We also need to subtract 1 in the very end since an empty chain is not allowed.

So how does solve() work? The method solveSimple() is reponsible for counting the total number of chains, forgetting about the symmetry issues. In order to calculate the number of symmetric chains, we use the following observation: a symmetric chain of length n is uniquely defined by its first half. When n is even, the half is n/2 long and each color appears twice less frequently in the half. When n is odd, the half is (n+1)/2 long and each color except the one of the middle jewel appears twice less frequently, while the number of occurrences of the color of the middle jewel is (old+1)/2. The code in lines 129-139 is basically concerned with the color of the middle jewel. When both #R and #G are odd, this means the middle jewel is both red and green, which is impossible. When only one of them is odd, this should be the color of the middle jewel, so the number of blue jewels should be even and we divide maxB by 2 to get the corresponding number for the half-chain. When both of them are even, either there's no middle jewel (line 133), or it is blue (line 136).

That's it! Questions? Comments?