<s> from sympy . core . compatibility import range , zip_longest <EOL> from sympy . utilities . enumerative import ( <EOL> list_visitor , <EOL> MultisetPartitionTraverser , <EOL> multiset_partitions_taocp <EOL> ) <EOL> from sympy . utilities . iterables import _set_partitions <EOL> from sympy . utilities . pytest import slow <EOL> def part_range_filter ( partition_iterator , lb , ub ) : <EOL> """<STR_LIT>""" <EOL> for state in partition_iterator : <EOL> f , lpart , pstack = state <EOL> if lpart >= lb and lpart < ub : <EOL> yield state <EOL> def multiset_partitions_baseline ( multiplicities , components ) : <EOL> """<STR_LIT>""" <EOL> canon = [ ] <EOL> for ct , elem in zip ( multiplicities , components ) : <EOL> canon . extend ( [ elem ] * ct ) <EOL> cache = set ( ) <EOL> n = len ( canon ) <EOL> for nc , q in _set_partitions ( n ) : <EOL> rv = [ [ ] for i in range ( nc ) ] <EOL> for i in range ( n ) : <EOL> rv [ q [ i ] ] . append ( canon [ i ] ) <EOL> canonical = tuple ( <EOL> sorted ( [ tuple ( p ) for p in rv ] ) ) <EOL> cache . add ( canonical ) <EOL> return cache <EOL> def compare_multiset_w_baseline ( multiplicities ) : <EOL> """<STR_LIT>""" <EOL> letters = "<STR_LIT>" <EOL> bl_partitions = multiset_partitions_baseline ( multiplicities , letters ) <EOL> aocp_partitions = set ( ) <EOL> for state in multiset_partitions_taocp ( multiplicities ) : <EOL> p1 = tuple ( sorted ( <EOL> [ tuple ( p ) for p in list_visitor ( state , letters ) ] ) ) <EOL> aocp_partitions . add ( p1 ) <EOL> assert bl_partitions == aocp_partitions <EOL> def compare_multiset_states ( s1 , s2 ) : <EOL> """<STR_LIT>""" <EOL> f1 , lpart1 , pstack1 = s1 <EOL> f2 , lpart2 , pstack2 = s2 <EOL> if ( lpart1 == lpart2 ) and ( f1 [ <NUM_LIT:0> : lpart1 + <NUM_LIT:1> ] == f2 [ <NUM_LIT:0> : lpart2 + <NUM_LIT:1> ] ) : <EOL> if pstack1 [ <NUM_LIT:0> : f1 [ lpart1 + <NUM_LIT:1> ] ] == pstack2 [ <NUM_LIT:0> : f2 [ lpart2 + <NUM_LIT:1> ] ] : <EOL> return True <EOL> return False <EOL> def test_multiset_partitions_taocp ( ) : <EOL> """<STR_LIT>""" <EOL> multiplicities = [ <NUM_LIT:2> , <NUM_LIT:2> ] <EOL> compare_multiset_w_baseline ( multiplicities ) <EOL> multiplicities = [ <NUM_LIT:4> , <NUM_LIT:3> , <NUM_LIT:1> ] <EOL> compare_multiset_w_baseline ( multiplicities ) <EOL> def test_multiset_partitions_versions ( ) : <EOL> """<STR_LIT>""" <EOL> multiplicities = [ <NUM_LIT:5> , <NUM_LIT:2> , <NUM_LIT:2> , <NUM_LIT:1> ] <EOL> m = MultisetPartitionTraverser ( ) <EOL> for s1 , s2 in zip_longest ( m . enum_all ( multiplicities ) , <EOL> multiset_partitions_taocp ( multiplicities ) ) : <EOL> assert compare_multiset_states ( s1 , s2 ) <EOL> def subrange_exercise ( mult , lb , ub ) : <EOL> """<STR_LIT>""" <EOL> m = MultisetPartitionTraverser ( ) <EOL> assert m . count_partitions ( mult ) == m . count_partitions_slow ( mult ) <EOL> ma = MultisetPartitionTraverser ( ) <EOL> mc = MultisetPartitionTraverser ( ) <EOL> md = MultisetPartitionTraverser ( ) <EOL> a_it = ma . enum_range ( mult , lb , ub ) <EOL> b_it = part_range_filter ( multiset_partitions_taocp ( mult ) , lb , ub ) <EOL> c_it = part_range_filter ( mc . enum_small ( mult , ub ) , lb , sum ( mult ) ) <EOL> d_it = part_range_filter ( md . enum_large ( mult , lb ) , <NUM_LIT:0> , ub ) <EOL> for sa , sb , sc , sd in zip_longest ( a_it , b_it , c_it , d_it ) : <EOL> assert compare_multiset_states ( sa , sb ) <EOL> assert compare_multiset_states ( sa , sc ) <EOL> assert compare_multiset_states ( sa , sd ) <EOL> def test_subrange ( ) : <EOL> mult = [ <NUM_LIT:4> , <NUM_LIT:4> , <NUM_LIT:2> , <NUM_LIT:1> ] <EOL> lb = <NUM_LIT:1> <EOL> ub = <NUM_LIT:2> <EOL> subrange_exercise ( mult , lb , ub ) <EOL> @ slow <EOL> def test_subrange_large ( ) : <EOL> mult = [ <NUM_LIT:6> , <NUM_LIT:3> , <NUM_LIT:2> , <NUM_LIT:1> ] <EOL> lb = <NUM_LIT:4> <EOL> ub = <NUM_LIT:7> <EOL> subrange_exercise ( mult , lb , ub ) </s>
