Guest User

Untitled

a guest
Mar 18th, 2013
90
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 7.73 KB | None | 0 0
  1. #!/usr/bin/env python3
  2. # -*- coding: utf-8 -*-
  3.  
  4.  
  5. def read_ints():
  6.     return [int(i) for i in input().split()]
  7.  
  8.  
  9. class MultidimensionalFenwickSumTree:
  10.     '''Build a multidimensional table of elements (usually numbers), providing
  11.    item assignment and multidimensional range sum queries.
  12.  
  13.    Time complexity of operations:
  14.      assume n is maximal index among all dimensions,
  15.             d is number of dimensions.
  16.      * construction: O(n^d log^d n)
  17.      * update: O(log^d n)
  18.      * range sum query: O(2^d log^d n)
  19.  
  20.    >>> mfst = MultidimensionalFenwickSumTree([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 0, 1]])
  21.    >>> mfst
  22.    MultidimensionalFenwickSumTree([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 0, 1]])
  23.    >>> mfst[0:2][1:3].sum()
  24.    14
  25.    >>> mfst[2][2]
  26.    0
  27.    >>> mfst[2][2] = 8
  28.    >>> mfst
  29.    MultidimensionalFenwickSumTree([[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 8, 1]])
  30.    >>> mfst[1:3][0:3].sum()
  31.    40
  32.    >>> mfst[-1][0:4].sum()
  33.    26
  34.    '''
  35.  
  36.     import copy
  37.  
  38.     class SliceView:
  39.         '''Provide access to elements and subtables of a Fenwick tree.
  40.        self.indices is maintained as a tuple of slice object with step == None
  41.        and 0 <= indices[i].start < self.mfst.length[i],
  42.            indices[i].start < indices[i].stop <= self.mfst.length[i].
  43.        '''
  44.  
  45.         def __init__(self, mfst, indices, subscript=False):
  46.             self.mfst = mfst
  47.             self.indices = indices
  48.             self.subscript = subscript or all([isinstance(i, int) for i in indices])
  49.  
  50.         def __fenwick_rec_update(self, indices, difference, level, subtable):
  51.             k = indices[level].start
  52.             while k < self.mfst.length[level]:
  53.                 if level + 1 == self.mfst.dim:
  54.                     subtable[k] += difference
  55.                 else:
  56.                     self.__fenwick_rec_update(indices, difference, level + 1, subtable[k])
  57.                 k = k | (k + 1)
  58.  
  59.         def __getitem__(self, index, value=None):
  60.             level = len(self.indices)
  61.             assert level < self.mfst.dim, 'too many levels of indices'
  62.  
  63.             if isinstance(index, int):
  64.                 if index < 0:
  65.                     index += self.mfst.length[level]
  66.                 index = slice(index, index + 1)
  67.             else:
  68.                 self.subscript = False
  69.  
  70.             if (not 0 <= index.start < self.mfst.length[level] or
  71.                 not index.start < index.stop <= self.mfst.length[level]):
  72.                 raise IndexError('index out of range')
  73.  
  74.             indices = self.indices + (index,)
  75.  
  76.             if (level + 1 == self.mfst.dim and self.subscript):
  77.                 tmp = self.mfst.table
  78.                 for i in indices[:-1]:
  79.                     tmp = tmp[i.start]
  80.                 if value is None:
  81.                     return tmp[indices[-1].start]
  82.                 else:
  83.                     difference = value - tmp[indices[-1].start]
  84.                     tmp[indices[-1].start] = value
  85.  
  86.                     self.__fenwick_rec_update(indices, difference, 0, self.mfst.sum)
  87.             else:
  88.                 if value is None:
  89.                     return type(self)(self.mfst, indices, self.subscript)
  90.                 else:
  91.                     if level + 1 == self.mfst.dim:
  92.                         raise IndexError('cannot assign to a slice')
  93.                     else:
  94.                         raise IndexError('not enough levels of indices')
  95.  
  96.         def __setitem__(self, index, value):
  97.             self.__getitem__(index, value)
  98.  
  99.         def __rec_prefix_sum(self, res, indices, level, subtable):
  100.             k = indices[level] - 1
  101.             while k >= 0:
  102.                 if level + 1 == self.mfst.dim:
  103.                     res[0] += subtable[k]
  104.                 else:
  105.                     self.__rec_prefix_sum(res, indices, level + 1, subtable[k])
  106.                 k = (k & (k + 1)) - 1
  107.  
  108.         def prefix_sum(self, indices):
  109.             res = [self.mfst.scalar_type()]
  110.             self.__rec_prefix_sum(res, indices, 0, self.mfst.sum)
  111.             return res[0]
  112.  
  113.         def __rec_sum(self, res, indices, parity, level):
  114.             if level == self.mfst.dim:
  115.                 res[0] += parity * self.prefix_sum(indices)
  116.             else:
  117.                 indices.append(self.indices[level].start)
  118.                 self.__rec_sum(res, indices, -parity, level + 1)
  119.                 indices.pop()
  120.                 indices.append(self.indices[level].stop)
  121.                 self.__rec_sum(res, indices, parity, level + 1)
  122.                 indices.pop()
  123.  
  124.         def sum(self):
  125.             res = [self.mfst.scalar_type()]
  126.             self.__rec_sum(res, [], 1, 0)
  127.             return res[0]
  128.  
  129.     def __init__(self, table):
  130.         '''Build a table. Parameter 'table' should consist of nested lists.
  131.        Everything which is not a list inside 'tables' is treated as a scalar
  132.        value.
  133.        '''
  134.         table = self.copy.deepcopy(table)
  135.         self.length = []
  136.         mainstream_subtable = table
  137.         stack_for_length_check = [mainstream_subtable]
  138.  
  139.         while isinstance(mainstream_subtable, list):
  140.             self.length.append(len(mainstream_subtable))
  141.             for subtable in stack_for_length_check:
  142.                 if len(subtable) != self.length[-1]:
  143.                     raise ValueError('length mismatch on level {0}: '
  144.                         'the subtable {1} should be of length {2} '
  145.                         'as the subtable {3} is'.format(level, subtable,
  146.                             self.length[-1], mainstream_subtable))
  147.             mainstream_subtable = mainstream_subtable[0]
  148.             stack_for_length_check = [j for i in stack_for_length_check
  149.                                         for j in i]
  150.  
  151.         self.scalar_type = type(mainstream_subtable)
  152.         self.dim = len(self.length)
  153.  
  154.         assert self.dim > 0
  155.         assert 0 not in self.length
  156.  
  157.         self.sum = [self.scalar_type() for i in range(self.length[-1])]
  158.         for i in self.length[-2::-1]:
  159.             self.sum = [self.copy.deepcopy(self.sum) for j in range(i)]
  160.  
  161.         self.table = self.copy.deepcopy(self.sum)
  162.  
  163.         def recursive_construction(level, slice_indices, subtable):
  164.             if level < self.dim:
  165.                 for i in range(self.length[level]):
  166.                     slice_indices.append(slice(i, i + 1))
  167.                     recursive_construction(level + 1, slice_indices, subtable[i])
  168.                     slice_indices.pop()
  169.             else:
  170.                 self.SliceView(self, tuple(slice_indices[:-1]), True)[slice_indices[-1].start] = subtable
  171.  
  172.         recursive_construction(0, [], table)
  173.  
  174.     def __repr__(self):
  175.         return 'MultidimensionalFenwickSumTree({0})'.format(repr(self.table))
  176.  
  177.     def __getitem__(self, index):
  178.         return self.SliceView(self, ())[index]
  179.  
  180.     def __setitem__(self, index, value):
  181.         self.SliceView(self, ())[index] = value
  182.  
  183.  
  184.  
  185. def main():
  186.     global input
  187.     global print
  188.    
  189.     #fin = open('input.txt', 'r')
  190.     #input = fin.readline
  191.     #fout = open('output.txt', 'w')
  192.     #_print = print
  193.     #print = lambda *args, **kwargs: _print(*args, file=fout, **kwargs)
  194.  
  195.     n = read_ints()[0]
  196.     mfst = MultidimensionalFenwickSumTree([[[0] * n for i in range(n)]
  197.                                                     for j in range(n)])
  198.     while True:
  199.         line = read_ints()
  200.         if line[0] == 1:
  201.             x, y, z, k = line[1:]
  202.             mfst[x][y][z] += k
  203.         elif line[0] == 2:
  204.             x1, y1, z1, x2, y2, z2 = line[1:]
  205.             print(mfst[x1:x2 + 1][y1:y2 + 1][z1:z2 + 1].sum())
  206.         else:
  207.             assert line[0] == 3
  208.             break
  209.  
  210.  
  211.     #fout.close()
  212.  
  213.  
  214. if __name__ == '__main__':
  215.     # import doctest
  216.     # doctest.testmod()
  217.     main()
Advertisement
Add Comment
Please, Sign In to add comment