<s> from __future__ import print_function , division <EOL> from sympy . core import Basic , Dict , sympify <EOL> from sympy . core . compatibility import as_int , default_sort_key , range <EOL> from sympy . functions . combinatorial . numbers import bell <EOL> from sympy . matrices import zeros <EOL> from sympy . sets . sets import FiniteSet <EOL> from sympy . utilities . iterables import has_dups , flatten , group <EOL> from collections import defaultdict <EOL> class Partition ( FiniteSet ) : <EOL> """<STR_LIT>""" <EOL> _rank = None <EOL> _partition = None <EOL> def __new__ ( cls , * partition ) : <EOL> """<STR_LIT>""" <EOL> args = partition <EOL> if not all ( isinstance ( part , ( list , FiniteSet ) ) for part in args ) : <EOL> raise ValueError ( <EOL> "<STR_LIT>" ) <EOL> partition = sorted ( sum ( partition , [ ] ) , key = default_sort_key ) <EOL> if has_dups ( partition ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> obj = FiniteSet . __new__ ( cls , * [ FiniteSet ( * x ) for x in args ] ) <EOL> obj . members = tuple ( partition ) <EOL> obj . size = len ( partition ) <EOL> return obj <EOL> def sort_key ( self , order = None ) : <EOL> """<STR_LIT>""" <EOL> if order is None : <EOL> members = self . members <EOL> else : <EOL> members = tuple ( sorted ( self . members , <EOL> key = lambda w : default_sort_key ( w , order ) ) ) <EOL> return list ( map ( default_sort_key , ( self . size , members , self . rank ) ) ) <EOL> @ property <EOL> def partition ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _partition is None : <EOL> self . _partition = sorted ( [ sorted ( p , key = default_sort_key ) <EOL> for p in self . args ] ) <EOL> return self . _partition <EOL> def __add__ ( self , other ) : <EOL> """<STR_LIT>""" <EOL> other = as_int ( other ) <EOL> offset = self . rank + other <EOL> result = RGS_unrank ( ( offset ) % <EOL> RGS_enum ( self . size ) , <EOL> self . size ) <EOL> return Partition . from_rgs ( result , self . members ) <EOL> def __sub__ ( self , other ) : <EOL> """<STR_LIT>""" <EOL> return self . __add__ ( - other ) <EOL> def __le__ ( self , other ) : <EOL> """<STR_LIT>""" <EOL> return self . sort_key ( ) <= sympify ( other ) . sort_key ( ) <EOL> def __lt__ ( self , other ) : <EOL> """<STR_LIT>""" <EOL> return self . sort_key ( ) < sympify ( other ) . sort_key ( ) <EOL> @ property <EOL> def rank ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _rank is not None : <EOL> return self . _rank <EOL> self . _rank = RGS_rank ( self . RGS ) <EOL> return self . _rank <EOL> @ property <EOL> def RGS ( self ) : <EOL> """<STR_LIT>""" <EOL> rgs = { } <EOL> partition = self . partition <EOL> for i , part in enumerate ( partition ) : <EOL> for j in part : <EOL> rgs [ j ] = i <EOL> return tuple ( [ rgs [ i ] for i in sorted ( <EOL> [ i for p in partition for i in p ] , key = default_sort_key ) ] ) <EOL> @ classmethod <EOL> def from_rgs ( self , rgs , elements ) : <EOL> """<STR_LIT>""" <EOL> if len ( rgs ) != len ( elements ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> max_elem = max ( rgs ) + <NUM_LIT:1> <EOL> partition = [ [ ] for i in range ( max_elem ) ] <EOL> j = <NUM_LIT:0> <EOL> for i in rgs : <EOL> partition [ i ] . append ( elements [ j ] ) <EOL> j += <NUM_LIT:1> <EOL> if not all ( p for p in partition ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> return Partition ( * partition ) <EOL> class IntegerPartition ( Basic ) : <EOL> """<STR_LIT>""" <EOL> _dict = None <EOL> _keys = None <EOL> def __new__ ( cls , partition , integer = None ) : <EOL> """<STR_LIT>""" <EOL> if integer is not None : <EOL> integer , partition = partition , integer <EOL> if isinstance ( partition , ( dict , Dict ) ) : <EOL> _ = [ ] <EOL> for k , v in sorted ( list ( partition . items ( ) ) , reverse = True ) : <EOL> if not v : <EOL> continue <EOL> k , v = as_int ( k ) , as_int ( v ) <EOL> _ . extend ( [ k ] * v ) <EOL> partition = tuple ( _ ) <EOL> else : <EOL> partition = tuple ( sorted ( map ( as_int , partition ) , reverse = True ) ) <EOL> sum_ok = False <EOL> if integer is None : <EOL> integer = sum ( partition ) <EOL> sum_ok = True <EOL> else : <EOL> integer = as_int ( integer ) <EOL> if not sum_ok and sum ( partition ) != integer : <EOL> raise ValueError ( "<STR_LIT>" % integer ) <EOL> if any ( i < <NUM_LIT:1> for i in partition ) : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> obj = Basic . __new__ ( cls , integer , partition ) <EOL> obj . partition = list ( partition ) <EOL> obj . integer = integer <EOL> return obj <EOL> def prev_lex ( self ) : <EOL> """<STR_LIT>""" <EOL> d = defaultdict ( int ) <EOL> d . update ( self . as_dict ( ) ) <EOL> keys = self . _keys <EOL> if keys == [ <NUM_LIT:1> ] : <EOL> return IntegerPartition ( { self . integer : <NUM_LIT:1> } ) <EOL> if keys [ - <NUM_LIT:1> ] != <NUM_LIT:1> : <EOL> d [ keys [ - <NUM_LIT:1> ] ] -= <NUM_LIT:1> <EOL> if keys [ - <NUM_LIT:1> ] == <NUM_LIT:2> : <EOL> d [ <NUM_LIT:1> ] = <NUM_LIT:2> <EOL> else : <EOL> d [ keys [ - <NUM_LIT:1> ] - <NUM_LIT:1> ] = d [ <NUM_LIT:1> ] = <NUM_LIT:1> <EOL> else : <EOL> d [ keys [ - <NUM_LIT:2> ] ] -= <NUM_LIT:1> <EOL> left = d [ <NUM_LIT:1> ] + keys [ - <NUM_LIT:2> ] <EOL> new = keys [ - <NUM_LIT:2> ] <EOL> d [ <NUM_LIT:1> ] = <NUM_LIT:0> <EOL> while left : <EOL> new -= <NUM_LIT:1> <EOL> if left - new >= <NUM_LIT:0> : <EOL> d [ new ] += left // new <EOL> left -= d [ new ] * new <EOL> return IntegerPartition ( self . integer , d ) <EOL> def next_lex ( self ) : <EOL> """<STR_LIT>""" <EOL> d = defaultdict ( int ) <EOL> d . update ( self . as_dict ( ) ) <EOL> key = self . _keys <EOL> a = key [ - <NUM_LIT:1> ] <EOL> if a == self . integer : <EOL> d . clear ( ) <EOL> d [ <NUM_LIT:1> ] = self . integer <EOL> elif a == <NUM_LIT:1> : <EOL> if d [ a ] > <NUM_LIT:1> : <EOL> d [ a + <NUM_LIT:1> ] += <NUM_LIT:1> <EOL> d [ a ] -= <NUM_LIT:2> <EOL> else : <EOL> b = key [ - <NUM_LIT:2> ] <EOL> d [ b + <NUM_LIT:1> ] += <NUM_LIT:1> <EOL> d [ <NUM_LIT:1> ] = ( d [ b ] - <NUM_LIT:1> ) * b <EOL> d [ b ] = <NUM_LIT:0> <EOL> else : <EOL> if d [ a ] > <NUM_LIT:1> : <EOL> if len ( key ) == <NUM_LIT:1> : <EOL> d . clear ( ) <EOL> d [ a + <NUM_LIT:1> ] = <NUM_LIT:1> <EOL> d [ <NUM_LIT:1> ] = self . integer - a - <NUM_LIT:1> <EOL> else : <EOL> a1 = a + <NUM_LIT:1> <EOL> d [ a1 ] += <NUM_LIT:1> <EOL> d [ <NUM_LIT:1> ] = d [ a ] * a - a1 <EOL> d [ a ] = <NUM_LIT:0> <EOL> else : <EOL> b = key [ - <NUM_LIT:2> ] <EOL> b1 = b + <NUM_LIT:1> <EOL> d [ b1 ] += <NUM_LIT:1> <EOL> need = d [ b ] * b + d [ a ] * a - b1 <EOL> d [ a ] = d [ b ] = <NUM_LIT:0> <EOL> d [ <NUM_LIT:1> ] = need <EOL> return IntegerPartition ( self . integer , d ) <EOL> def as_dict ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _dict is None : <EOL> groups = group ( self . partition , multiple = False ) <EOL> self . _keys = [ g [ <NUM_LIT:0> ] for g in groups ] <EOL> self . _dict = dict ( groups ) <EOL> return self . _dict <EOL> @ property <EOL> def conjugate ( self ) : <EOL> """<STR_LIT>""" <EOL> j = <NUM_LIT:1> <EOL> temp_arr = list ( self . partition ) + [ <NUM_LIT:0> ] <EOL> k = temp_arr [ <NUM_LIT:0> ] <EOL> b = [ <NUM_LIT:0> ] * k <EOL> while k > <NUM_LIT:0> : <EOL> while k > temp_arr [ j ] : <EOL> b [ k - <NUM_LIT:1> ] = j <EOL> k -= <NUM_LIT:1> <EOL> j += <NUM_LIT:1> <EOL> return b <EOL> def __lt__ ( self , other ) : <EOL> """<STR_LIT>""" <EOL> return list ( reversed ( self . partition ) ) < list ( reversed ( other . partition ) ) <EOL> def __le__ ( self , other ) : <EOL> """<STR_LIT>""" <EOL> return list ( reversed ( self . partition ) ) <= list ( reversed ( other . partition ) ) <EOL> def as_ferrers ( self , char = '<STR_LIT:#>' ) : <EOL> """<STR_LIT>""" <EOL> return "<STR_LIT:\n>" . join ( [ char * i for i in self . partition ] ) <EOL> def __str__ ( self ) : <EOL> return str ( list ( self . partition ) ) <EOL> def random_integer_partition ( n , seed = None ) : <EOL> """<STR_LIT>""" <EOL> from sympy . utilities . randtest import _randint <EOL> n = as_int ( n ) <EOL> if n < <NUM_LIT:1> : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> randint = _randint ( seed ) <EOL> partition = [ ] <EOL> while ( n > <NUM_LIT:0> ) : <EOL> k = randint ( <NUM_LIT:1> , n ) <EOL> mult = randint ( <NUM_LIT:1> , n // k ) <EOL> partition . append ( ( k , mult ) ) <EOL> n -= k * mult <EOL> partition . sort ( reverse = True ) <EOL> partition = flatten ( [ [ k ] * m for k , m in partition ] ) <EOL> return partition <EOL> def RGS_generalized ( m ) : <EOL> """<STR_LIT>""" <EOL> d = zeros ( m + <NUM_LIT:1> ) <EOL> for i in range ( <NUM_LIT:0> , m + <NUM_LIT:1> ) : <EOL> d [ <NUM_LIT:0> , i ] = <NUM_LIT:1> <EOL> for i in range ( <NUM_LIT:1> , m + <NUM_LIT:1> ) : <EOL> for j in range ( m ) : <EOL> if j <= m - i : <EOL> d [ i , j ] = j * d [ i - <NUM_LIT:1> , j ] + d [ i - <NUM_LIT:1> , j + <NUM_LIT:1> ] <EOL> else : <EOL> d [ i , j ] = <NUM_LIT:0> <EOL> return d <EOL> def RGS_enum ( m ) : <EOL> """<STR_LIT>""" <EOL> if ( m < <NUM_LIT:1> ) : <EOL> return <NUM_LIT:0> <EOL> elif ( m == <NUM_LIT:1> ) : <EOL> return <NUM_LIT:1> <EOL> else : <EOL> return bell ( m ) <EOL> def RGS_unrank ( rank , m ) : <EOL> """<STR_LIT>""" <EOL> if m < <NUM_LIT:1> : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> if rank < <NUM_LIT:0> or RGS_enum ( m ) <= rank : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> L = [ <NUM_LIT:1> ] * ( m + <NUM_LIT:1> ) <EOL> j = <NUM_LIT:1> <EOL> D = RGS_generalized ( m ) <EOL> for i in range ( <NUM_LIT:2> , m + <NUM_LIT:1> ) : <EOL> v = D [ m - i , j ] <EOL> cr = j * v <EOL> if cr <= rank : <EOL> L [ i ] = j + <NUM_LIT:1> <EOL> rank -= cr <EOL> j += <NUM_LIT:1> <EOL> else : <EOL> L [ i ] = int ( rank / v + <NUM_LIT:1> ) <EOL> rank %= v <EOL> return [ x - <NUM_LIT:1> for x in L [ <NUM_LIT:1> : ] ] <EOL> def RGS_rank ( rgs ) : <EOL> """<STR_LIT>""" <EOL> rgs_size = len ( rgs ) <EOL> rank = <NUM_LIT:0> <EOL> D = RGS_generalized ( rgs_size ) <EOL> for i in range ( <NUM_LIT:1> , rgs_size ) : <EOL> n = len ( rgs [ ( i + <NUM_LIT:1> ) : ] ) <EOL> m = max ( rgs [ <NUM_LIT:0> : i ] ) <EOL> rank += D [ n , m + <NUM_LIT:1> ] * rgs [ i ] <EOL> return rank </s>
