UVA 11997 - K Smallest Sums ( Priority Queue, Merge )

UVa Online Judge
First I thought I could simply make a structure to record where each index of the arrays are, and expand them in priority queue, however, did not work out because for example when ( 0, 1 ) and ( 1, 0 ) are both visited, ( 1, 1 ) will be visited twice. Anyways, we will come up to that the number of arrays can be reduced, as a bunch of subproblems, merging each pair of arrays until there is only one left. This problem is a good example for the approach of simplifying the original problem to find the relation.

typedef pair< int, int > pii;
typedef vector< int > vi;
typedef vector< vi > vvi;

vi merge( const vi &a, const vi &b ){
    vi res;
    priority_queue< pii, vector< pii >, greater< pii > > pq;
    for( int i = 0; i < a.size(); ++i )
        pq.push( pii( a[ i ] + b[ 0 ], 0 ) );
    while( res.size() < a.size() ){
        int sum = pq.top().first;
        int bid = pq.top().second;
        pq.pop();
        res.push_back( sum );
        if( bid + 1 < b.size() )
            pq.push( pii( sum - b[ bid ] + b[ bid + 1 ], bid + 1 ) );
    }
    return res;
}

void solve( int N ){
    vvi A( N, vi( N ) );
    for( int i = 0; i < N; ++i )
        for( int j = 0; j < N; ++j )
            cin >> A[ i ][ j ];

    for( int i = 0; i < N; ++i )
        sort( A[ i ].begin(), A[ i ].end() );

    vi cur = A[ 0 ];
    for( int i = 1; i < N; ++i )
        cur = merge( cur, A[ i ] );

    for( int i = 0; i < N; ++i )
        cout << cur[ i ] << ( i + 1 == N ? '\n' : ' ' );
}