User role based on conditional logic

Ive found this snippet that sets the role based on email, is there a way to tweak it so change based on “company” selected?

add_action(“um_registration_complete”,”um_042121_change_user_role”, 1, 2 );
function um_042121_change_user_role( $user_id, $args ){

if( isset( $args['user_email'] ) ){
list($user, $domain) = explode('@', $args['user_email'] );
if ($domain == 'email.co.uk') {
wp_update_user( array( 'ID' => $user_id, 'role' => 'um_staff' ) );
}
}

thanks

• This reply was modified 2 hours, 45 minutes ago by damiansmith86.

Tweaked the code myself and seems to be working as expected… just incase anyone else needs this:

add_action("um_registration_complete","um_042121_change_user_role", 1, 2 );
function um_042121_change_user_role( $user_id, $args ){
if ( isset( $args['company'] ) && strstr( $args['company'], 'Company Name' ) ) {
wp_update_user( array( 'ID' => $user_id, 'role' => 'um_staff' ) );
}
}