got the digonals down

1 files changed, 34 insertions, 5 deletions

diff --git a/2010/1a/a.py b/2010/1a/a.py
index 1e253ff..47b3fe7 100644
--- a/2010/1a/a.py
+++ b/2010/1a/a.py
@@ -1,6 +1,24 @@
 f = open('1.in', 'r')
 o = open('1.out', 'w')
 
+def cols(m):
+    n = len(m)
+    return [''.join([m[j][i] for j in range(n)]) for i in range(n)]
+
+def diag(m):
+    n = len(m)
+    d = [zip(range(i+1) ,range(i+1)[::-1]) for i in range(n)]
+    print d
+    return [''.join([m[i][j] for (i,j) in diag]) for diag in d]
+    
+def diag_inv(m):
+    n = len(m)
+    d = [zip(range(i+1),range(i+1)[::-1]) for i in range(n)]
+    return [''.join([m[n-1-i][j] for (i,j) in diag]) for diag in d]
+
+print diag(['abc','def','ghi'])
+exit()
+
 def rotate(m):
     n = len(m)
     return [''.join([m[j][i] for j in range(n)[::-1]]) for i in range(n)]
@@ -10,13 +28,24 @@ def droprow(r):
 
 def gravity(m):
     return map(droprow, m)
-    
-def win(m):
+
+def winrow(r, k):
+    res = 0
+    if r.find('B'*k) != -1:
+        res |= 1
+    if r.find('R'*k) != -1:
+        res |= 2
+    return res
+
+def win(m, k):
+    R = ['Neither', 'Blue', 'Red', 'Both']
     win = 0
-    for (i, c) in ((1, 'B'), (2, 'R')):
-        win |= i
+    for row in m:
+        win |= winrow(row, k)
+    for col in cols(m):
+        win |= winrow(col, k)
     
-    return ['Neither', 'Blue', 'Red', 'Both'][win]
+    return R[win]
 
 t = int(f.readline().strip())